Add arch_svg generator source (meta-device introspection + modular diff + SVG render)

arch_svg/README.md

@@ -0,0 +1,170 @@
+# `arch_svg` — architecture diagrams generated from the `transformers` source
+
+Introspect any model in `transformers` and emit its architecture as an **SVG**, in the
+spirit of Sebastian Raschka's [LLM Architecture Gallery](https://sebastianraschka.com/llm-architecture-gallery/)
+— but generated automatically from the library's own source instead of drawn by hand.
+
+Output is a **static, self-contained SVG** (CSS only, no JS, no external fetch) — meant to be
+embeddable directly in a model's README on the Hub. Modular models get **two** views
+(`--mode both`, the default for `all`); standalone models get only the full view.
+
+- **full** — a detailed, Sebastian-Raschka-style diagram read **top-to-bottom like the code**:
+  `input_ids [1, 12]` → token embedding → the **expanded decoder block** (pre-norm →
+  **RoPE** → self-attention with Q/K/V/O projections at real meta-tensor dims, GQA/MQA/MLA,
+  QK-norm → ⊕ residual → pre-norm → MLP gate/up/down **or** Sparse MoE router→top-k→experts
+  (+shared)) → final norm → LM head → softmax → `logits`, with residual skip arrows and
+  tensor shapes at each stage. A left **layer-schedule strip** shows the per-layer
+  `config.layer_types` pattern (e.g. sliding/full, linear/full), and per-distinct-type
+  **attention-mask pattern grids** are drawn on the right. The facts panel (far right) shows
+  the example checkpoint id, model_type, dims, heads/kv, experts, etc.
+- **diff** — *only what this model changes* vs. the parent(s) it inherits from in its
+  `modular_<name>.py`. A clean modular model produces a tiny, legible diff; a bloated one
+  produces a large diff. **The diff size is a direct, automatic measure of how modular each
+  model actually is.** A "changes by class" panel lists every overridden / added / deleted
+  member.
+
+> Compared to interactive viewers like hfviewer.com, this is deliberately **static**: one
+> deterministic `.svg` per model, safe to commit and embed as a model-card fingerprint.
+
+The headline feature is **exploiting modular `transformers`**: many models are defined as a
+`modular_<name>.py` that inherits from another model (`class GemmaModel(LlamaModel)`), and
+the diff view shows exactly the override/add/delete payload of that inheritance.
+
+No network, no weights. Models are built on `torch.device("meta")` from their **default
+config class** (`Config()` — no `config.json` download required), so the whole zoo renders
+in seconds at zero GPU/memory cost. Anything that can't be built on meta falls back to
+config-only parsing; anything that can't be parsed at all is recorded in `report.json`
+rather than crashing the batch.
+
+## Usage
+
+```bash
+python -m arch_svg one  --model llama  --mode full --out out/   # detailed Raschka-style block
+python -m arch_svg one  --model gemma  --mode diff --out out/   # Gemma's diff vs Llama
+python -m arch_svg all  --mode both    --out out/ --jobs 10      # whole library, BOTH views + index.html
+python -m arch_svg all  --mode both    --out out/ --limit 30     # quick subset
+python -m arch_svg list                                          # list discovered models
+```
+
+`all` (mode `both`) writes `out/<model>.svg` (full) **and** `out/<model>.diff.svg` (diff) per
+model, an `out/index.html` contact sheet (each card shows both thumbnails, with a live filter
+box and a "largest diffs" leaderboard), and `out/report.json` with per-model status
+(`ok` / `built-from-config-only` / `failed` + traceback), the resolved parent, and the count
+of overridden / added / deleted / new nodes.
+
+Open `out/index.html` in a browser — SVGs use CSS variables and honor `prefers-color-scheme`,
+so light/dark is a one-line swap. (Note: `librsvg`/Quick Look don't support CSS custom
+properties, so a *browser* is the intended viewer; opening a single `.svg` in some image
+apps may show it un-themed.)
+
+## How it works
+
+```
+arch_svg/
+  discover.py    # enumerate models/* (import-free; stat the tree + the config mapping)
+  modular.py     # ← the interesting part: libcst diff of modular_*.py vs its parent(s)
+  introspect.py  # default-config + meta-device build → normalized ArchModel (+ detectors)
+  layout.py      # ArchModel/ArchDiff → positioned boxes (standardized vocabulary)
+  render.py      # boxes → SVG string (themeable <style>, deterministic output)
+  gallery.py     # batch driver, index.html, report.json (keeps going on failure)
+  cli.py         # python -m arch_svg
+```
+
+`modular.py` does **not** re-run `utils/modular_model_converter.py` (that exists to
+*generate* standalone modeling files). Instead it mirrors the converter's inheritance
+semantics with `libcst` to recover the diff payload directly:
+
+- a method/attr present in **both** the modular class and its named parent → **overridden**
+- present **only** in the modular class → **added**
+- a deletion sentinel (`attr = AttributeError(...)` or `raise AttributeError`) → **deleted**
+
+The parent model is read from the modular file's `from ..<model>.modeling_<model> import …`
+imports — the same import-following the converter performs.
+
+### Standardized visual vocabulary
+
+Following the "**standardize, don't abstract**" tenet, every component kind maps to a fixed
+shape/color, so the same block looks identical across all 450+ diagrams:
+
+| color | component | | color | component |
+|---|---|---|---|---|
+| 🔵 blue | token embedding | | 🟠 orange | MoE / experts |
+| 🟦 cyan | attention | | 🟣 purple | MLP / dense FFN |
+| 🟢 green | Mamba / SSM | | ⚪ grey | norm (RMS/Layer) |
+| 🔴 red | LM head | | 🟡 yellow | RoPE |
+
+In **diff mode** the parent architecture is drawn ghosted, and only changed regions are
+highlighted: **green = added, amber = overridden, red = deleted**. A "changes by class"
+panel enumerates every modular class and exactly which methods/attrs it touches.
+
+## What diff mode revealed about modularity
+
+Run over the full library (454 models with modeling code): **249 modular, 205 standalone**;
+374 built on meta, 72 config-only, 8 unrenderable.
+
+**Llama is the spine of the library.** 46 models inherit (transitively, via their dominant
+parent) from `llama` — then `vit` (10), `llava` (9), `mixtral` (9), `clip` (7),
+`wav2vec2` (7). A handful of base models account for most of the zoo.
+
+**The cleanest modular models are nearly free.** Diff "size" =
+`overridden + added + deleted + 3·new_classes`. Median modular diff is **39**; the
+exemplars are tiny:
+
+| model | parent | diff size |
+|---|---|---|
+| `glm` | llama | 2 |
+| `qwen3` | qwen2 | 4 |
+| `qwen2` | llama | 5 |
+| `granite` | llama | 7 |
+| `ijepa` | vit | 7 |
+
+9 models are **pure-override** (only re-tune existing methods, add *nothing* and create *no*
+new classes) — the gold standard for modular: `bitnet`, `camembert`, `ernie4_5`, `glm`,
+`granite`, `ijepa`, `qwen2`, `qwen3`, `xlm_roberta`.
+
+**The largest diffs flag where modularity breaks down (refactor candidates).** 30 models
+have diff size > 100:
+
+| model | parent | size | reading |
+|---|---|---|---|
+| `qwen2_5_omni` | qwen2_5_vl | 429 | omni model bolts on audio+talker towers — mostly *new* classes |
+| `qwen3_omni_moe` | qwen3_moe | 310 | same pattern |
+| `gemma3n` | gemma3 | 213 | adds a full audio encoder + per-layer laurel/altup |
+| `rt_detr` | detr | 197 | **0 overridden, 137 added** — inherits the *name* but rebuilds the model |
+| `maskformer` | detr | 143 | **0 overridden, 86 added** — same: near-zero reuse of the parent |
+
+Two distinct shapes of "large diff" show up, and the diagram tells them apart at a glance:
+
+1. **Legitimately large** — multimodal/omni models (`qwen2_5_omni`, `phi4_multimodal`,
+   `florence2`) that genuinely add vision/audio sub-towers. Big, but the additions are real
+   new architecture.
+2. **Suspicious** — models with **0 overridden + many added** (`rt_detr`, `maskformer`).
+   They declare a parent but override none of its behaviour; the inheritance buys almost
+   nothing and is closer to copy-paste-with-a-base-class. These are the clearest candidates
+   to either lean harder on the parent or stop pretending to inherit from it.
+
+### A note on philosophy (and why there are 8 failures + 72 config-only)
+
+This tool deliberately contains **no per-model escapes or special-cases**. The resolvers
+(model-type normalization, class-name→region mapping, sequence-valued config formatting) are
+all generic. When a model *doesn't* render cleanly, that is treated as a **finding, not a bug
+to paper over**:
+
+- the **8 failures** are all composite meta-models (`encoder_decoder`, `rag`, `musicgen`,
+  `vision_encoder_decoder`, …) that can't be instantiated from defaults;
+- the **72 config-only** models are ones whose default config can't build on meta;
+- a **huge diff** means the modeling code deviates from the standard.
+
+The right response to a model that doesn't fit the standard vocabulary is to **standardize the
+model**, not to add a workaround here. In that sense this gallery doubles as a linter for
+modularity and standardization across the library.
+
+## Regenerating the gallery
+
+```bash
+python -m arch_svg all --mode diff --out arch_svg/out      --jobs 10
+python -m arch_svg all --mode full --out arch_svg/out_full --jobs 10
+```
+
+Output is deterministic (stable ordering, integer coordinates) so the SVGs are diff-able in
+git.

arch_svg/__init__.py

@@ -0,0 +1,23 @@
+"""arch_svg - introspect any `transformers` model and emit an architecture diagram as SVG.
+
+The headline feature is exploiting *modular* transformers: a model defined as a
+``modular_<name>.py`` that inherits from another model can be rendered in ``diff`` mode,
+showing only what it changes relative to its parent(s). A clean modular model produces a
+tiny, legible diff -- a direct measure of how modular the library really is.
+
+See ``python -m arch_svg --help``.
+"""
+
+from .introspect import ArchModel, LayerBlock, introspect
+from .modular import ArchDiff, ClassChange, diff, resolve_parent
+
+
+__all__ = [
+    "ArchModel",
+    "LayerBlock",
+    "introspect",
+    "ArchDiff",
+    "ClassChange",
+    "diff",
+    "resolve_parent",
+]

arch_svg/__main__.py

@@ -0,0 +1,7 @@
+import sys
+
+from .cli import main
+
+
+if __name__ == "__main__":
+    sys.exit(main())

arch_svg/cli.py

@@ -0,0 +1,84 @@
+"""`python -m arch_svg ...`
+
+python -m arch_svg one --model llama --mode full --out out/
+python -m arch_svg one --model gemma --mode diff --out out/
+python -m arch_svg all --mode diff --out out/ --jobs 8
+python -m arch_svg all --mode diff --out out/ --limit 20   # quick subset
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+
+
+def _cmd_one(args) -> int:
+    from .discover import discover_models, model_type_for
+    from .gallery import render_one
+
+    model_type = None
+    for e in discover_models():
+        if e.name == args.model:
+            model_type = model_type_for(e)
+            break
+    svg, record = render_one(args.model, model_type or args.model, args.mode)
+    os.makedirs(args.out, exist_ok=True)
+    path = os.path.join(args.out, f"{args.model}.svg")
+    with open(path, "w", encoding="utf-8") as f:
+        f.write(svg)
+    print(f"wrote {path}  (status={record['status']})")
+    if record.get("build_error"):
+        print(f"  note: {record['build_error']}")
+    if args.mode == "diff" and record.get("is_modular") is not None:
+        print(
+            f"  modular={record.get('is_modular')} parent={record.get('parent_model')} totals={record.get('diff_totals')}"
+        )
+    return 0 if record["status"] != "failed" else 1
+
+
+def _cmd_all(args) -> int:
+    from .gallery import run_all
+
+    res = run_all(out=args.out, mode=args.mode, jobs=args.jobs, limit=args.limit)
+    s = res["summary"]
+    print(f"\ndone → {os.path.join(args.out, 'index.html')}")
+    print(f"  {s}")
+    return 0
+
+
+def _cmd_list(args) -> int:
+    from .discover import discover_models
+
+    for e in discover_models():
+        tag = "modular" if e.has_modular else "standalone"
+        print(f"{e.name:32s} {tag}")
+    return 0
+
+
+def main(argv=None) -> int:
+    p = argparse.ArgumentParser(prog="arch_svg", description="transformers architecture SVG generator")
+    sub = p.add_subparsers(dest="cmd", required=True)
+
+    one = sub.add_parser("one", help="render a single model")
+    one.add_argument("--model", required=True)
+    one.add_argument("--mode", choices=["full", "diff"], default="full")
+    one.add_argument("--out", default="out")
+    one.set_defaults(func=_cmd_one)
+
+    alle = sub.add_parser("all", help="render the whole model zoo + index.html")
+    alle.add_argument("--mode", choices=["full", "diff", "both"], default="both")
+    alle.add_argument("--out", default="out")
+    alle.add_argument("--jobs", type=int, default=4)
+    alle.add_argument("--limit", type=int, default=None)
+    alle.set_defaults(func=_cmd_all)
+
+    lst = sub.add_parser("list", help="list discovered models")
+    lst.set_defaults(func=_cmd_list)
+
+    args = p.parse_args(argv)
+    return args.func(args)
+
+
+if __name__ == "__main__":
+    sys.exit(main())

arch_svg/discover.py

@@ -0,0 +1,113 @@
+"""Enumerate models under ``transformers/models/*`` -- no imports of the models themselves.
+
+Discovery is cheap and import-free: we only stat the directory tree and read the
+``model_type`` from the config mapping. Heavy work (building the model on meta) happens
+later, per-model, so a single broken model can never crash the whole batch.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import os
+from dataclasses import dataclass
+from functools import lru_cache
+
+
+# Subpackages under models/ that are not actual models (utilities, tokenizers-only, etc.).
+_NON_MODELS = {"auto", "__pycache__"}
+
+
+@dataclass
+class ModelEntry:
+    """A discovered model subpackage. Purely path-derived; nothing is imported yet."""
+
+    name: str  # directory name, e.g. "gemma"
+    path: str  # absolute path to the model subpackage
+    modeling_files: list[str]  # modeling_*.py present
+    modular_file: str | None  # modular_*.py if present
+    config_files: list[str]  # configuration_*.py present
+
+    @property
+    def has_modular(self) -> bool:
+        return self.modular_file is not None
+
+
+@lru_cache(maxsize=1)
+def models_root() -> str:
+    """Locate ``transformers/models`` inside the *installed* source tree."""
+    spec = importlib.util.find_spec("transformers")
+    if spec is None or not spec.submodule_search_locations:
+        raise RuntimeError("Could not locate the installed `transformers` package.")
+    return os.path.join(list(spec.submodule_search_locations)[0], "models")
+
+
+def discover_models(root: str | None = None) -> list[ModelEntry]:
+    """Return every model subpackage under ``transformers/models/`` sorted by name."""
+    root = root or models_root()
+    entries: list[ModelEntry] = []
+    for name in sorted(os.listdir(root)):
+        if name in _NON_MODELS or name.startswith("_") or name.startswith("."):
+            continue
+        path = os.path.join(root, name)
+        if not os.path.isdir(path):
+            continue
+        files = os.listdir(path)
+        modeling = sorted(f for f in files if f.startswith("modeling_") and f.endswith(".py"))
+        configs = sorted(f for f in files if f.startswith("configuration_") and f.endswith(".py"))
+        modular = next((f for f in sorted(files) if f.startswith("modular_") and f.endswith(".py")), None)
+        if not modeling and not modular and not configs:
+            continue  # tokenizer/processor-only subpackage
+        entries.append(
+            ModelEntry(
+                name=name,
+                path=path,
+                modeling_files=modeling,
+                modular_file=modular,
+                config_files=configs,
+            )
+        )
+    return entries
+
+
+@lru_cache(maxsize=1)
+def _config_mapping_keys() -> dict[str, str]:
+    """Map model directory name heuristics to model_type via the config mapping.
+
+    Returns a dict {model_type: model_type}; the keys ARE the model_types registered in
+    the library. Kept import-light: only imports the auto config module.
+    """
+    from transformers.models.auto.configuration_auto import CONFIG_MAPPING_NAMES
+
+    return dict(CONFIG_MAPPING_NAMES)
+
+
+def model_type_for(entry: ModelEntry) -> str | None:
+    """Best-effort resolution of a directory name to a registered ``model_type``.
+
+    The directory name usually equals the model_type; when it does not (e.g. dir
+    ``kosmos2`` -> ``kosmos-2``, ``x_clip`` -> ``xclip``), fall back to normalized matching
+    against the config mapping (strip ``-``/``_`` and compare). For dirs that back several
+    model_types (``data2vec`` -> ``data2vec-{audio,text,vision}``) the first is chosen.
+    """
+    keys = list(_config_mapping_keys())
+    name = entry.name
+    if name in keys:
+        return name
+    # exact match after hyphen/underscore unification
+    for mt in keys:
+        if mt.replace("-", "_") == name:
+            return mt
+
+    def norm(s: str) -> str:
+        return s.replace("-", "").replace("_", "")
+
+    target = norm(name)
+    candidates = [mt for mt in keys if norm(mt) == target]
+    if not candidates:
+        # dir name is a prefix of exactly one (or more) model_types, e.g. lasr -> lasr_ctc
+        candidates = [mt for mt in keys if norm(mt).startswith(target)]
+    if candidates:
+        # avoid obviously auxiliary types; prefer the shortest / alphabetically first
+        candidates = [c for c in candidates if "privacy" not in c and "filter" not in c] or candidates
+        return sorted(candidates, key=lambda c: (len(c), c))[0]
+    return None

arch_svg/engine.py

@@ -0,0 +1,71 @@
+"""Tiny self-contained layout primitives — our own engine, no external deps.
+
+The point of this module is **text-aware sizing**: every other part of the package can ask
+"how wide is this string at this font size?" and "make this string fit in this width", so
+boxes are sized to their content and text never overflows or runs together. Kept dependency-
+free and deterministic (a fixed proportional-width table) so output stays reproducible.
+"""
+
+from __future__ import annotations
+
+
+# per-character width as a fraction of the font size, for a generic sans-serif.
+# Approximate but stable; good enough to size/ellipsize boxes so text fits.
+_NARROW = set("iIl.,:;'!|()[]{}ftj ")
+_WIDE = set("mMW@%—–→⊙⊕")
+_MID = set("rtszc")
+
+
+def char_em(ch: str) -> float:
+    if ch in _NARROW:
+        return 0.30
+    if ch in _WIDE:
+        return 0.92
+    if ch in _MID:
+        return 0.48
+    if ch.isupper() or ch.isdigit():
+        return 0.62
+    return 0.54
+
+
+def text_width(s: str, px: float) -> float:
+    """Estimated rendered width (in px) of ``s`` at font size ``px``."""
+    return sum(char_em(c) for c in s) * px
+
+
+def fit_text(s: str, max_px: float, px: float) -> str:
+    """Return ``s`` unchanged if it fits in ``max_px``, else truncated with a trailing ellipsis."""
+    if max_px <= 0 or text_width(s, px) <= max_px:
+        return s
+    ell = "…"
+    budget = max_px - text_width(ell, px)
+    out = []
+    w = 0.0
+    for c in s:
+        cw = char_em(c) * px
+        if w + cw > budget:
+            break
+        out.append(c)
+        w += cw
+    return ("".join(out).rstrip() + ell) if out else ell
+
+
+def wrap_text(s: str, max_px: float, px: float, max_lines: int = 2) -> list[str]:
+    """Greedily wrap ``s`` to ``max_px`` over up to ``max_lines`` lines (last line ellipsized)."""
+    if text_width(s, px) <= max_px:
+        return [s]
+    words = s.split(" ")
+    lines: list[str] = []
+    cur = ""
+    for w in words:
+        trial = (cur + " " + w).strip()
+        if text_width(trial, px) <= max_px or not cur:
+            cur = trial
+        else:
+            lines.append(cur)
+            cur = w
+            if len(lines) == max_lines - 1:
+                break
+    rest = " ".join(words[sum(len(line.split(" ")) for line in lines) :]) if lines else cur
+    lines.append(fit_text(rest or cur, max_px, px))
+    return lines[:max_lines]

arch_svg/gallery.py

@@ -0,0 +1,250 @@
+"""Batch driver: render every model, write per-model SVGs, an index.html contact sheet,
+and a report.json. Keeps going on per-model failure -- errors are collected, never fatal.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import traceback
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from html import escape
+
+from .discover import ModelEntry, discover_models, model_type_for
+from .introspect import introspect
+from .layout import build_diff, build_full
+from .modular import diff
+from .render import render
+
+
+def render_one(model: str, model_type: str | None, mode: str) -> tuple[str, dict]:
+    """Render a single model to an SVG string + a report record. Never raises."""
+    record: dict = {"model": model, "mode": mode, "status": "ok"}
+    try:
+        am = introspect(model, model_type)
+        record["build_status"] = am.status
+        record["decoder_kind"] = am.decoder_kind
+        record["layer_summary"] = am.layer_summary
+        if am.error:
+            record["build_error"] = am.error
+        if mode in ("diff", "both"):
+            ad = diff(model)
+            record["is_modular"] = ad.is_modular
+            record["parent_model"] = ad.parent_model
+            record["parent_models"] = ad.parent_models
+            record["diff_totals"] = ad.totals
+            record["note"] = ad.note
+            svg = render(build_diff(am, ad))
+        else:
+            svg = render(build_full(am))
+        if am.status == "failed":
+            record["status"] = "failed"
+        elif am.status == "config-only":
+            record["status"] = "built-from-config-only"
+        return svg, record
+    except Exception:
+        record["status"] = "failed"
+        record["traceback"] = traceback.format_exc()[-1500:]
+        # emit a minimal error placeholder svg so the gallery link still works
+        svg = (
+            f'<svg xmlns="http://www.w3.org/2000/svg" width="480" height="120" viewBox="0 0 480 120">'
+            f'<rect width="480" height="120" fill="#fff"/>'
+            f'<text x="16" y="40" font-family="sans-serif" font-size="16" fill="#dc2626">{escape(model)} — failed</text>'
+            f'<text x="16" y="66" font-family="monospace" font-size="11" fill="#6b7280">{escape(mode)} mode</text>'
+            f"</svg>"
+        )
+        return svg, record
+
+
+def _write(out, fname, svg):
+    path = os.path.join(out, fname)
+    with open(path, "w", encoding="utf-8") as f:
+        f.write(svg)
+    return os.path.basename(path)
+
+
+def _worker(args):
+    model, model_type, mode, out = args
+    if mode == "both":
+        from .modular import _modular_path
+
+        # the detailed full view is always emitted
+        full_svg, record = render_one(model, model_type, "full")
+        record["svg"] = _write(out, f"{model}.svg", full_svg)
+        # the diff view is only meaningful for modular models; standalone -> full only
+        if _modular_path(model) is not None:
+            diff_svg, diff_rec = render_one(model, model_type, "diff")
+            record["diff_svg"] = _write(out, f"{model}.diff.svg", diff_svg)
+            for key in ("is_modular", "parent_model", "parent_models", "diff_totals", "note"):
+                if key in diff_rec:
+                    record[key] = diff_rec[key]
+            if diff_rec.get("status") == "failed" and record.get("status") != "failed":
+                record["diff_status"] = "failed"
+        else:
+            record["is_modular"] = False
+            record["parent_model"] = None
+            record["note"] = "standalone (no modular parent)"
+        return record
+    svg, record = render_one(model, model_type, mode)
+    record["svg"] = _write(out, f"{model}.svg", svg)
+    return record
+
+
+def run_all(out: str, mode: str = "both", jobs: int = 4, limit: int | None = None) -> dict:
+    os.makedirs(out, exist_ok=True)
+    entries: list[ModelEntry] = discover_models()
+    if limit:
+        entries = entries[:limit]
+    tasks = [(e.name, model_type_for(e), mode, out) for e in entries]
+
+    records: list[dict] = []
+    if jobs and jobs > 1:
+        with ProcessPoolExecutor(max_workers=jobs) as ex:
+            futures = {ex.submit(_worker, t): t[0] for t in tasks}
+            for i, fut in enumerate(as_completed(futures), 1):
+                name = futures[fut]
+                try:
+                    records.append(fut.result())
+                except Exception:
+                    records.append(
+                        {
+                            "model": name,
+                            "mode": mode,
+                            "status": "failed",
+                            "traceback": traceback.format_exc()[-800:],
+                            "svg": f"{name}.svg",
+                        }
+                    )
+                print(f"[{i}/{len(tasks)}] {name}", flush=True)
+    else:
+        for i, t in enumerate(tasks, 1):
+            records.append(_worker(t))
+            print(f"[{i}/{len(tasks)}] {t[0]}", flush=True)
+
+    records.sort(key=lambda r: r["model"])
+    report = _summarize(records, mode)
+    with open(os.path.join(out, "report.json"), "w", encoding="utf-8") as f:
+        json.dump({"mode": mode, "summary": report, "models": records}, f, indent=2)
+    with open(os.path.join(out, "index.html"), "w", encoding="utf-8") as f:
+        f.write(_index_html(records, mode, report))
+    return {"summary": report, "out": out}
+
+
+def _summarize(records: list[dict], mode: str) -> dict:
+    s = {"total": len(records), "ok": 0, "config_only": 0, "failed": 0}
+    if mode in ("diff", "both"):
+        s["modular"] = 0
+        s["standalone"] = 0
+        s["largest_diffs"] = []
+    for r in records:
+        st = r.get("status")
+        if st == "failed":
+            s["failed"] += 1
+        elif st == "built-from-config-only":
+            s["config_only"] += 1
+        else:
+            s["ok"] += 1
+        if mode in ("diff", "both"):
+            if r.get("is_modular"):
+                s["modular"] += 1
+            elif r.get("is_modular") is False:
+                s["standalone"] += 1
+    if mode in ("diff", "both"):
+        scored = []
+        for r in records:
+            t = r.get("diff_totals")
+            if r.get("is_modular") and t:
+                size = t["overridden"] + t["added"] + t["deleted"] + 3 * t["new_classes"]
+                scored.append((size, r["model"], r.get("parent_model"), t))
+        scored.sort(reverse=True)
+        s["largest_diffs"] = [{"model": m, "parent": p, "diff_size": sz, "totals": t} for sz, m, p, t in scored[:20]]
+    return s
+
+
+def _index_html(records: list[dict], mode: str, summary: dict) -> str:
+    cards = []
+    for r in records:
+        st = r.get("status", "ok")
+        badge = {"ok": "#16a34a", "built-from-config-only": "#d97706", "failed": "#dc2626"}.get(st, "#6b7280")
+        sub = ""
+        if mode in ("diff", "both"):
+            if r.get("is_modular"):
+                t = r.get("diff_totals", {})
+                sub = f"↳ {escape(str(r.get('parent_model')))} · ovr {t.get('overridden', 0)} / add {t.get('added', 0)} / del {t.get('deleted', 0)} / new {t.get('new_classes', 0)}"
+            else:
+                sub = "standalone"
+        else:
+            sub = escape(str(r.get("layer_summary") or r.get("decoder_kind") or ""))
+        if mode == "both" and r.get("diff_svg"):
+            thumbs = (
+                f'<div class="thumbs">'
+                f'<a href="{escape(r["svg"])}" target="_blank" class="tw">'
+                f'<object data="{escape(r["svg"])}" type="image/svg+xml" class="thumb"></object>'
+                f'<span class="tlabel">full</span></a>'
+                f'<a href="{escape(r["diff_svg"])}" target="_blank" class="tw">'
+                f'<object data="{escape(r["diff_svg"])}" type="image/svg+xml" class="thumb"></object>'
+                f'<span class="tlabel">diff</span></a></div>'
+            )
+        else:
+            thumbs = (
+                f'<a href="{escape(r["svg"])}" target="_blank" class="tw">'
+                f'<object data="{escape(r["svg"])}" type="image/svg+xml" class="thumb"></object></a>'
+            )
+        cards.append(
+            f'<div class="card">'
+            f'<div class="card-h"><span class="name">{escape(r["model"])}</span>'
+            f'<span class="dot" style="background:{badge}"></span></div>'
+            f"{thumbs}"
+            f'<div class="sub">{sub}</div></div>'
+        )
+
+    summary_html = f"ok {summary['ok']} · config-only {summary['config_only']} · failed {summary['failed']}"
+    if mode in ("diff", "both"):
+        summary_html += f" · modular {summary.get('modular', 0)} · standalone {summary.get('standalone', 0)}"
+        rows = "".join(
+            f"<tr><td>{escape(d['model'])}</td><td>{escape(str(d['parent']))}</td>"
+            f"<td>{d['diff_size']}</td><td>{d['totals']['overridden']}</td>"
+            f"<td>{d['totals']['added']}</td><td>{d['totals']['deleted']}</td>"
+            f"<td>{d['totals']['new_classes']}</td></tr>"
+            for d in summary.get("largest_diffs", [])
+        )
+        big = (
+            "<h2>Largest diffs (least modular)</h2>"
+            '<table class="lb"><tr><th>model</th><th>parent</th><th>size</th>'
+            "<th>ovr</th><th>add</th><th>del</th><th>new</th></tr>"
+            f"{rows}</table>"
+        )
+    else:
+        big = ""
+
+    return f"""<!doctype html>
+<html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1">
+<title>transformers architecture gallery — {escape(mode)}</title>
+<style>
+:root {{ color-scheme: light dark; --bg:#fff; --fg:#111; --muted:#666; --card:#f6f8fa; --bd:#e5e7eb; }}
+@media (prefers-color-scheme: dark) {{ :root {{ --bg:#0d1117; --fg:#e6edf3; --muted:#8b949e; --card:#161b22; --bd:#30363d; }} }}
+body {{ background:var(--bg); color:var(--fg); font-family:ui-sans-serif,-apple-system,"Segoe UI",Roboto,sans-serif; margin:0; padding:28px; }}
+h1 {{ margin:0 0 4px; }} .meta {{ color:var(--muted); margin-bottom:18px; }}
+.grid {{ display:grid; grid-template-columns:repeat(auto-fill,minmax(340px,1fr)); gap:16px; }}
+.card {{ background:var(--card); border:1px solid var(--bd); border-radius:12px; padding:12px; color:inherit; display:flex; flex-direction:column; }}
+.card:hover {{ border-color:#3b82f6; }}
+.card-h {{ display:flex; justify-content:space-between; align-items:center; }}
+.name {{ font-weight:700; }} .dot {{ width:10px; height:10px; border-radius:50%; }}
+.thumbs {{ display:flex; gap:8px; }}
+.tw {{ position:relative; flex:1; text-decoration:none; }}
+.tlabel {{ position:absolute; top:10px; left:6px; font-size:10px; font-weight:700; color:var(--muted); background:var(--bg); padding:1px 5px; border-radius:4px; opacity:.85; }}
+.thumb {{ width:100%; height:260px; pointer-events:none; margin:8px 0; background:var(--bg); border-radius:6px; border:1px solid var(--bd); }}
+.sub {{ font-size:12px; color:var(--muted); }}
+table.lb {{ border-collapse:collapse; margin:8px 0 24px; font-size:13px; }}
+table.lb td, table.lb th {{ border:1px solid var(--bd); padding:4px 10px; text-align:left; }}
+input {{ padding:8px 12px; border-radius:8px; border:1px solid var(--bd); background:var(--card); color:var(--fg); width:260px; margin-bottom:16px; }}
+</style></head>
+<body>
+<h1>transformers architecture gallery</h1>
+<div class="meta">mode: <b>{escape(mode)}</b> · {summary_html}</div>
+{big}
+<input id="q" placeholder="filter models…" oninput="for(const c of document.querySelectorAll('.card')){{c.style.display=c.querySelector('.name').textContent.includes(this.value)?'':'none'}}">
+<div class="grid">
+{"".join(cards)}
+</div>
+</body></html>"""

arch_svg/introspect.py

@@ -0,0 +1,1402 @@
+"""Build a normalized ``ArchModel`` from a model's default config + meta-device module tree.
+
+No network, no weights: we instantiate the *config class* with its defaults and build the
+model under ``torch.device("meta")`` so allocation is free. If the meta build fails we fall
+back to a config-only ``ArchModel`` (status ``"config-only"``). A registry of detectors maps
+module class names to normalized component kinds so unknown modules degrade to their class
+name rather than crashing.
+"""
+
+from __future__ import annotations
+
+import re
+import warnings
+from dataclasses import dataclass, field
+
+
+warnings.filterwarnings("ignore")
+
+
+# --------------------------------------------------------------------------------- model
+
+
+@dataclass
+class LayerBlock:
+    """A run of consecutive identical decoder layers, collapsed into one box."""
+
+    layer_class: str
+    count: int
+    kind: str  # "attention" | "moe" | "mamba" | "linear_attention" | "recurrent" | "other"
+    attn_variant: str | None = None  # MHA | GQA | MQA | MLA | sliding | None
+    mlp_kind: str = "dense"  # dense | moe
+    norm: str | None = None
+
+
+# ---- structured "inside the block" specs (for the Raschka-style expanded full view) ----
+
+
+@dataclass
+class Proj:
+    """A single projection (Linear) inside attention / MLP, with meta-tensor shapes."""
+
+    name: str
+    in_f: int | None
+    out_f: int | None
+    bias: bool = False
+
+
+@dataclass
+class AttentionSpec:
+    cls: str
+    variant: str | None = None  # MHA | GQA | MQA | MLA
+    n_heads: int | None = None
+    n_kv: int | None = None
+    head_dim: int | None = None
+    projs: list[Proj] = field(default_factory=list)
+    qk_norm: bool = False
+    rope: bool = False
+    sliding_window: int | None = None
+
+
+@dataclass
+class MLPSpec:
+    cls: str
+    is_moe: bool = False
+    act: str = "SiLU"
+    gate: Proj | None = None
+    up: Proj | None = None
+    down: Proj | None = None
+    n_experts: int | None = None
+    top_k: int | None = None
+    expert_dim: int | None = None
+    n_shared: int | None = None
+    projs: list[Proj] = field(default_factory=list)  # generic fallback when gate/up/down naming differs
+
+
+@dataclass
+class BlockSpec:
+    """The canonical pre-norm transformer block, decomposed for detailed rendering."""
+
+    layer_class: str
+    pre_attn_norm: str | None = None
+    attention: AttentionSpec | None = None
+    mixer: AttentionSpec | None = None  # token mixer for non-attention blocks (Mamba/SSM/RWKV/linear-attn)
+    post_attn_norm: str | None = None
+    mlp: MLPSpec | None = None
+    extra: list[str] = field(default_factory=list)  # e.g. sandwich-norm note
+
+
+@dataclass
+class ArchModel:
+    model: str
+    model_type: str | None = None
+    config_class: str | None = None
+    status: str = "ok"  # ok | config-only | failed
+    error: str | None = None
+
+    # scalar config facts
+    hidden_size: int | None = None
+    vocab_size: int | None = None
+    num_layers: int | None = None
+    num_attention_heads: int | None = None
+    num_kv_heads: int | None = None
+    head_dim: int | None = None
+    intermediate_size: int | None = None
+    max_position_embeddings: int | None = None
+    sliding_window: int | None = None
+
+    # normalized architecture facts
+    norm_type: str | None = None
+    positional: str | None = None
+    attn_variant: str | None = None
+    decoder_kind: str = "dense"  # dense | sparse-moe | hybrid | ssm | other
+    is_moe: bool = False
+    num_experts: int | None = None
+    experts_per_token: int | None = None
+    tie_word_embeddings: bool | None = None
+
+    layer_blocks: list[LayerBlock] = field(default_factory=list)
+    layer_summary: str | None = None  # e.g. "36× DeltaNet + 12× Attention"
+    module_kinds: dict[str, int] = field(default_factory=dict)  # class name -> count (top kinds)
+    top_class: str | None = None  # outer model class name
+    block: BlockSpec | None = None  # decomposed decoder block internals (full view)
+    generic_tree: list[dict] = field(default_factory=list)  # full module tree for non-transformer models
+    flow: list[dict] = field(default_factory=list)  # per-stage forward I/O shapes (data-flow view)
+    flow_input: str | None = None  # input tensor label, e.g. "input_values [1, 1, 8000]"
+    flow_output: list[int] | None = None  # final output shape
+    is_pipeline: bool = False  # ≥3 distinct top-level stage sub-networks (e.g. BLT, codecs)
+    rope_theta: float | None = None
+    hidden_act: str | None = None
+    layer_types: list[str] = field(default_factory=list)  # config.layer_types schedule (per layer)
+    alt_blocks: list[str] = field(default_factory=list)  # other layer classes (hybrid) not shown in detail
+    checkpoint: str | None = None  # example HF checkpoint id (from @auto_docstring), offline
+    attn_patterns: dict = field(default_factory=dict)  # layer_type -> 0/1 mask grid
+    tokens: list[str] = field(default_factory=list)  # example tokenized input shown at the top
+    # sparse/compressed attention components (e.g. DeepSeek-V4 HCA/CSA), per distinct layer_type
+    sparse_components: list[dict] = field(default_factory=list)
+    # fully-decomposed module trees (recursive) of each distinct attention variant + the FFN
+    attention_variants: list[dict] = field(default_factory=list)
+    mlp_tree: dict | None = None
+    mlp_variants: list[dict] = field(default_factory=list)  # per distinct mlp_layer_type (moe/hash_moe/dense)
+    mlp_layer_types: list[str] = field(default_factory=list)
+    # multimodal (VLM / audio) structure
+    is_multimodal: bool = False
+    towers: list[dict] = field(default_factory=list)  # encoder towers (vision/audio) + projector
+    cross_attention_layers: list | None = None  # LLM layer indices that cross-attend to features
+    modal_inputs: list[str] = field(default_factory=list)  # e.g. ["input_ids", "pixel_values"]
+    auto_classes: list[str] = field(default_factory=list)  # Auto* classes used to build sub-models
+    image_bidirectional: bool = False  # VLM: image tokens attend bidirectionally (prefix-LM)
+    family: str | None = None  # causal_lm | image_classification | audio_classification | seq2seq | ...
+    head_class: str | None = None  # the task wrapper class, e.g. ViTForImageClassification
+    view: str = "decoder"  # decoder | encoder | enc_dec | multimodal — which layout to use
+    num_labels: int | None = None
+
+
+# ----------------------------------------------------------------------------- detectors
+
+_MOE_RE = re.compile(r"(expert|router|moe|sparsemoe|sparse_moe)", re.IGNORECASE)
+_MAMBA_RE = re.compile(r"(mamba|mixer|ssm|\bs4\b|selectivescan)", re.IGNORECASE)
+_LINEAR_ATTN_RE = re.compile(r"(deltanet|linearattention|lineattn|gateddelta|gla\b|retention|rwkv)", re.IGNORECASE)
+_ATTN_RE = re.compile(r"attention$", re.IGNORECASE)
+_NORM_RE = re.compile(r"norm", re.IGNORECASE)
+
+
+def _kind_for_layer(class_name: str) -> str:
+    n = class_name.lower()
+    if _MAMBA_RE.search(n):
+        return "mamba"
+    if _LINEAR_ATTN_RE.search(n):
+        return "linear_attention"
+    if "recurrent" in n or "rnn" in n:
+        return "recurrent"
+    return "attention"
+
+
+def layer_type_kind(t: str) -> str:
+    """Classify a ``config.layer_types`` entry (e.g. 'sliding_attention') into a kind.
+
+    Sliding/full/compressed are all *attention* (they differ only in the mask); linear /
+    delta / gated / mamba / recurrent are structurally different layers.
+    """
+    n = t.lower()
+    if _MAMBA_RE.search(n):
+        return "mamba"
+    if _LINEAR_ATTN_RE.search(n) or "linear" in n or "delta" in n or "gated" in n:
+        return "linear_attention"
+    if "recurrent" in n:
+        return "recurrent"
+    return "attention"
+
+
+def _cfg(config, *names, default=None):
+    for n in names:
+        if hasattr(config, n):
+            v = getattr(config, n)
+            if v is not None:
+                return v
+    return default
+
+
+def _attn_variant(config) -> str | None:
+    heads = _cfg(config, "num_attention_heads", "n_head", "num_heads")
+    kv = _cfg(config, "num_key_value_heads", "num_kv_heads", default=heads)
+    # Multi-head Latent Attention (DeepSeek style)
+    if _cfg(config, "kv_lora_rank", "q_lora_rank") is not None:
+        return "MLA"
+    if heads is None:
+        return None
+    if kv is None:
+        return "MHA"
+    if kv == 1:
+        return "MQA"
+    if kv < heads:
+        return "GQA"
+    return "MHA"
+
+
+def _positional(config, module_classes: set[str]) -> str:
+    if any("rotary" in c.lower() or "rope" in c.lower() for c in module_classes):
+        return "RoPE"
+    if _cfg(config, "rope_parameters", "rope_theta", "rope_scaling") is not None:
+        return "RoPE"
+    if _cfg(config, "alibi") or any("alibi" in c.lower() for c in module_classes):
+        return "ALiBi"
+    if any(c in ("PositionEmbedding", "LearnedPositionalEmbedding") for c in module_classes):
+        return "learned-absolute"
+    # presence of a position embedding nn.Embedding sized to max_position_embeddings
+    return "RoPE" if _cfg(config, "rotary_emb_base") is not None else "n/a"
+
+
+def _norm_type(module_classes: set[str]) -> str | None:
+    norms = sorted({c for c in module_classes if _NORM_RE.search(c)})
+    if not norms:
+        return None
+    # prefer the model-specific norm name; normalize common ones
+    for c in norms:
+        if "rms" in c.lower():
+            return "RMSNorm"
+    for c in norms:
+        if "layernorm" in c.lower():
+            return "LayerNorm"
+    return norms[0]
+
+
+# ----------------------------------------------------------------------- module-tree walk
+
+
+def _find_layer_list(model, num_layers):
+    """Find the nn.ModuleList holding the decoder/encoder layers."""
+    import torch.nn as nn
+
+    best = None
+    for name, module in model.named_modules():
+        if isinstance(module, nn.ModuleList) and len(module) > 0:
+            # prefer one whose length matches num_layers, else the longest
+            if num_layers and len(module) == num_layers:
+                return name, module
+            if best is None or len(module) > len(best[1]):
+                best = (name, module)
+    return best if best else (None, None)
+
+
+def _collapse(class_names: list[str]) -> list[tuple[str, int]]:
+    out: list[tuple[str, int]] = []
+    for name in class_names:
+        if out and out[-1][0] == name:
+            out[-1] = (name, out[-1][1] + 1)
+        else:
+            out.append((name, 1))
+    return out
+
+
+def _all_configs(config):
+    """Yield the config and every nested sub-config (``*_config``)."""
+    yield config
+    for k in list(vars(config)):
+        if k.endswith("_config"):
+            sc = getattr(config, k, None)
+            if sc is not None and hasattr(sc, "model_type"):
+                yield from _all_configs(sc)
+
+
+def _sanitize_config(config) -> bool:
+    """Fix invalid default-config combos that break a meta build (generic, no per-model code).
+
+    The common one: ``hidden_size % num_attention_heads != 0`` in a sub-config (e.g. some
+    vision towers). We snap ``num_attention_heads`` to the largest divisor ≤ the original and
+    make ``num_key_value_heads`` divide it. Returns True if anything changed.
+    """
+    changed = False
+    for c in _all_configs(config):
+        h = getattr(c, "hidden_size", None)
+        nh = getattr(c, "num_attention_heads", None)
+        if isinstance(h, int) and isinstance(nh, int) and nh > 0 and h % nh != 0:
+            d = next((x for x in range(nh, 0, -1) if h % x == 0), 1)
+            c.num_attention_heads = d
+            changed = True
+            kv = getattr(c, "num_key_value_heads", None)
+            if isinstance(kv, int) and (kv == 0 or d % kv != 0):
+                c.num_key_value_heads = d
+    return changed
+
+
+def _build_meta_model(config, device: str = "meta"):
+    import torch
+
+    import transformers as tf
+
+    # try AutoModel first, then common task factories, then EVERY remaining AutoModelFor*
+    # (so depth / segmentation / time-series / retrieval wrappers all build without hardcoding)
+    preferred = [
+        "AutoModel",
+        "AutoModelForImageTextToText",
+        "AutoModelForCausalLM",
+        "AutoModelForImageClassification",
+        "AutoModelForDepthEstimation",
+        "AutoModelForSemanticSegmentation",
+        "AutoModelForPreTraining",
+        "AutoModelForSeq2SeqLM",
+    ]
+    rest = sorted(n for n in dir(tf) if n.startswith("AutoModel") and n not in preferred)
+    factories, seen = [], set()
+    for n in preferred + rest:
+        f = getattr(tf, n, None)
+        if f is not None and hasattr(f, "from_config") and n not in seen:
+            seen.add(n)
+            factories.append(f)
+    with torch.device(device):
+        last = None
+        for sanitize in (False, True):
+            if sanitize and not _sanitize_config(config):
+                break
+            for fac in factories:
+                try:
+                    return fac.from_config(config)
+                except Exception as e:
+                    last = e
+        raise last
+
+
+_CKPT_RE = re.compile(r"""checkpoint\s*=\s*["']([\w\-./]+/[\w\-.]+)["']""")
+
+
+def _example_checkpoint(model: str, mt: str) -> str | None:
+    """Best-effort example checkpoint id from ``@auto_docstring(checkpoint=...)`` in the
+    model's source (offline; we just read the files). Returns e.g. 'google/gemma-7b'."""
+    import os
+
+    from .discover import models_root
+
+    d = os.path.join(models_root(), model)
+    if not os.path.isdir(d):
+        return None
+    for f in sorted(os.listdir(d)):
+        if f.startswith(("modeling_", "configuration_", "modular_")) and f.endswith(".py"):
+            try:
+                m = _CKPT_RE.search(open(os.path.join(d, f), encoding="utf-8").read())
+            except Exception:
+                continue
+            if m:
+                return m.group(1)
+    return None
+
+
+_ACT_NAMES = {"gelu", "relu", "silu", "tanh", "sigmoid", "swish", "mish", "quickgelu", "newgelu", "geglu"}
+
+
+def _classify_mod(name: str, cls: str) -> str:
+    """Map a submodule to a *standardized* semantic kind used for colouring across every view."""
+    n = (name + " " + cls).lower()
+    c = cls.lower()
+    if "rotary" in n or "rope" in n:
+        return "rope"
+    if "norm" in c:
+        return "norm"
+    if "conv" in c:  # Conv1d/2d/3d, ParametrizedConv*, depthwise conv, ...
+        return "conv"
+    if any(k in c for k in ("lstm", "gru", "rnn")):
+        return "recurrent"
+    if "embedding" in c or "embed" in c:
+        return "embedding"
+    if "pool" in c:
+        return "pool"
+    if "dropout" in c or c == "identity" or "drop_path" in n:
+        return "dropout"
+    if "activation" in c or c in _ACT_NAMES or name in ("act", "activation", "act_fn"):
+        return "act"
+    if "quantiz" in c or "codebook" in c:
+        return "quantizer"
+    if "indexer" in n:
+        return "indexer"
+    if "compress" in n:
+        return "compressor"
+    if name == "gate" or "router" in n:
+        return "router"
+    if "expert" in n:
+        return "experts"
+    if any(k in n for k in ("classifier", "lm_head", "score")) or name == "head":
+        return "head"
+    if "linear" in c or "proj" in name:
+        return "linear"
+    return "other"
+
+
+def _decompose(module, depth: int) -> list[dict]:
+    """Recursively decompose a module into a tree of ``{name, cls, kind, dim, children}``."""
+    import torch.nn as nn
+
+    nodes = []
+    for n, c in module.named_children():
+        cls = type(c).__name__
+        kind = _classify_mod(n, cls)
+        dim = None
+        if isinstance(c, nn.Linear):
+            dim = f"[{c.in_features}→{c.out_features}]"
+        else:
+            w = getattr(c, "weight", None)
+            if w is not None and hasattr(w, "shape") and 1 <= len(w.shape) <= 3:
+                dim = "×".join(str(s) for s in tuple(w.shape))
+        is_leaf = isinstance(c, nn.Linear) or kind in ("norm", "rope")
+        children = _decompose(c, depth - 1) if (depth > 0 and not is_leaf) else []
+        nodes.append({"name": n, "cls": cls, "kind": kind, "dim": dim, "children": children})
+    return nodes
+
+
+def _mod_dim(c) -> str | None:
+    """A compact dimension label for any leaf module (Linear / Conv / Embedding / parametrized)."""
+    import torch.nn as nn
+
+    if isinstance(c, nn.Linear):
+        return f"[{c.in_features}→{c.out_features}]"
+    if isinstance(c, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
+        k = "×".join(str(x) for x in c.kernel_size)
+        g = f" g{c.groups}" if c.groups != 1 else ""
+        return f"{c.in_channels}→{c.out_channels} k{k}{g}"
+    if isinstance(c, nn.Embedding):
+        return f"[{c.num_embeddings}×{c.embedding_dim}]"
+    w = getattr(c, "weight", None)
+    if w is not None and hasattr(w, "shape") and 1 <= len(w.shape) <= 2:
+        return "×".join(str(s) for s in tuple(w.shape))
+    return None
+
+
+def _is_leaf_mod(c) -> bool:
+    import torch.nn as nn
+
+    return isinstance(c, (nn.Linear, nn.Conv1d, nn.Conv2d, nn.Conv3d, nn.Embedding)) or bool(
+        _NORM_RE.search(type(c).__name__)
+    )
+
+
+def _node_for(name, c, depth) -> dict:
+    leaf = _is_leaf_mod(c) or not list(c.named_children())
+    return {
+        "name": name,
+        "cls": type(c).__name__,
+        "kind": _classify_mod(name, type(c).__name__),
+        "dim": _mod_dim(c),
+        "children": _module_tree(c, depth - 1) if (depth > 0 and not leaf) else [],
+    }
+
+
+def _collapse_sequence(kids, depth: int) -> list[dict]:
+    """Collapse a sequence of children by detecting the smallest repeating *period*: a run like
+    ``(Conv1d, ResnetBlock) ×4`` becomes one ``×4`` group, ``[Stage, Stage, Stage] ×3`` becomes
+    ``×3 Stage``. Non-repeating elements (an input stem, a trailing LSTM, ...) stay as singles."""
+    sigs = [type(c).__name__ for _, c in kids]
+    n = len(kids)
+    out: list[dict] = []
+    i = 0
+    while i < n:
+        best = None
+        for p in range(1, (n - i) // 2 + 1):  # smallest period that repeats ≥ twice wins
+            k = 1
+            while i + (k + 1) * p <= n and sigs[i + k * p : i + (k + 1) * p] == sigs[i : i + p]:
+                k += 1
+            if k >= 2:
+                best = (p, k)
+                break
+        if best:
+            p, k = best
+            members = kids[i : i + p]
+            if p == 1:  # a homogeneous run → "×k ClassName" with the block's internals
+                nm, c = members[0]
+                g = _node_for(nm, c, depth)
+                g["name"], g["cls"] = f"[{i}:{i + k}]", f"{k}× {type(c).__name__}"
+                out.append(g)
+            else:  # a periodic pattern → "×k (A, B, …)" containing one copy of the pattern
+                out.append(
+                    {
+                        "name": f"[{i}:{i + p * k}]",
+                        "cls": f"{k}× (" + ", ".join(sigs[i : i + p]) + ")",
+                        "kind": "other",
+                        "dim": None,
+                        "children": [_node_for(nm, c, depth) for nm, c in members],
+                    }
+                )
+            i += p * k
+        else:
+            out.append(_node_for(*kids[i], depth))
+            i += 1
+    return out
+
+
+def _module_tree(module, depth: int) -> list[dict]:
+    """Recursive module hierarchy with repeating ``ModuleList``/``Sequential`` runs collapsed to
+    ``×N`` (see ``_collapse_sequence``). This renders *any* architecture — conv backbones, codecs,
+    FFT mixers, detectors — from its real structure, never a config-only schematic."""
+    import torch.nn as nn
+
+    nodes: list[dict] = []
+    for n, c in module.named_children():
+        cls = type(c).__name__
+        kids = list(c.named_children())
+        if isinstance(c, (nn.ModuleList, nn.Sequential)) and len(kids) > 1:
+            groups = _collapse_sequence(kids, depth) if depth > 0 else []
+            if len(groups) == 1 and groups[0]["cls"].startswith(f"{len(kids)}× "):
+                # whole list is one homogeneous repeat → inline it (e.g. "stages · 4× ResNetStage")
+                g = groups[0]
+                nodes.append({"name": n, "cls": g["cls"], "kind": g["kind"], "dim": None, "children": g["children"]})
+            else:
+                nodes.append(
+                    {
+                        "name": n,
+                        "cls": f"{cls} [{len(kids)}]",
+                        "kind": _classify_mod(n, cls),
+                        "dim": None,
+                        "children": groups,
+                    }
+                )
+            continue
+        nodes.append(_node_for(n, c, depth))
+    return nodes
+
+
+def _out_shape(x) -> list[int] | None:
+    """Best-effort tensor shape of a module's input/output: tensor, tuple/list, or ModelOutput."""
+    import torch
+
+    if torch.is_tensor(x):
+        return list(x.shape)
+    if isinstance(x, (list, tuple)):
+        for e in x:
+            s = _out_shape(e)
+            if s:
+                return s
+        return None
+    for attr in ("last_hidden_state", "logits", "sample", "audio_values", "audio_codes", "hidden_states"):
+        v = getattr(x, attr, None)
+        s = _out_shape(v) if v is not None else None
+        if s:
+            return s
+    return None
+
+
+def _dummy_inputs(model, config, device: str) -> dict:
+    """A minimal forward input matching the model's modality (no tokenizer / processor needed)."""
+    import torch
+
+    mi = getattr(model, "main_input_name", None) or "input_ids"
+    if mi in ("input_values", "input_features"):
+        return {mi: torch.zeros(1, 1, 8000, device=device)}
+    if mi == "pixel_values":
+        sz = getattr(config, "image_size", None) or 64
+        sz = sz if isinstance(sz, int) else (sz[0] if isinstance(sz, (list, tuple)) else 64)
+        ch = getattr(config, "num_channels", None) or 3
+        return {mi: torch.zeros(1, ch, sz, sz, device=device)}
+    return {"input_ids": torch.ones(1, 8, dtype=torch.long, device=device)}
+
+
+class _Deadline:
+    """Hard wall-clock timeout (SIGALRM) so a slow/looping forward can never hang the build.
+    Works in the main thread of each ProcessPoolExecutor worker; a no-op where SIGALRM is absent."""
+
+    def __init__(self, seconds: float):
+        self.seconds = seconds
+
+    def __enter__(self):
+        import signal
+
+        self._sig = getattr(signal, "SIGALRM", None)
+        if self._sig is not None:
+
+            def _raise(*_):
+                raise TimeoutError("forward exceeded time budget")
+
+            self._old = signal.signal(self._sig, _raise)
+            signal.setitimer(signal.ITIMER_REAL, self.seconds)
+        return self
+
+    def __exit__(self, *exc):
+        import signal
+
+        if self._sig is not None:
+            signal.setitimer(signal.ITIMER_REAL, 0)
+            signal.signal(self._sig, self._old)
+        return False
+
+
+def _run_flow(model, config, device: str) -> tuple[str, list[dict], list[int] | None] | None:
+    """Run one forward pass capturing the in→out shape of each top-level stage (in call order),
+    via forward hooks. Returns (input_label, stages, output_shape) or None on failure."""
+    import torch
+
+    model.eval()
+    inp = _dummy_inputs(model, config, device)
+    in_name = next(iter(inp))
+    in_shape = list(next(iter(inp.values())).shape)
+    flow: list[dict] = []
+    hooks = []
+    order = {id(m): i for i, (_, m) in enumerate(model.named_children())}
+
+    def mk(nm):
+        def hook(mod, i, o):
+            flow.append(
+                {
+                    "name": nm,
+                    "cls": type(mod).__name__,
+                    "in_shape": _out_shape(i),
+                    "out_shape": _out_shape(o),
+                    "_ord": order.get(id(mod), 999),
+                }
+            )
+
+        return hook
+
+    for nm, mod in model.named_children():
+        hooks.append(mod.register_forward_hook(mk(nm)))
+    out = None
+    try:
+        with torch.no_grad(), _Deadline(8):
+            out = model(**inp)
+    except BaseException:  # includes TimeoutError from the deadline
+        for h in hooks:
+            h.remove()
+        return None
+    for h in hooks:
+        h.remove()
+    if not flow:
+        return None
+    # de-dup (a stage hit more than once keeps its first call) and keep call order
+    seen, uniq = set(), []
+    for f in flow:
+        if f["name"] in seen:
+            continue
+        seen.add(f["name"])
+        uniq.append(f)
+    return (f"{in_name} {in_shape}", uniq, _out_shape(out))
+
+
+def _shrink_config(config):
+    """Shrink a config in-place so it can be instantiated cheaply on CPU for a real forward."""
+    for k in (
+        "num_hidden_layers",
+        "num_layers",
+        "decoder_layers",
+        "encoder_layers",
+        "n_layers",
+        "num_local_experts",
+        "n_routed_experts",
+    ):
+        if isinstance(getattr(config, k, None), int):
+            setattr(config, k, min(2, getattr(config, k)))
+    # cap the big allocations (vocab × hidden embeddings, position tables) so a CPU build is cheap
+    for k, cap in (("vocab_size", 256), ("max_position_embeddings", 64), ("max_target_positions", 64)):
+        if isinstance(getattr(config, k, None), int) and getattr(config, k) > cap:
+            setattr(config, k, cap)
+    for sub in ("encoder", "decoder", "text_config", "vision_config", "audio_config"):
+        c = getattr(config, sub, None)
+        if hasattr(c, "__dict__"):
+            _shrink_config(c)
+    return config
+
+
+def _stage_children(model_obj) -> list:
+    """Top-level children that are substantial *stage* sub-networks (own ModuleList or ≥2 kids),
+    i.e. a real pipeline rather than a single backbone + head. Excludes embeddings/norms/heads."""
+    import torch.nn as nn
+
+    stages = []
+    for n, c in model_obj.named_children():
+        if _is_leaf_mod(c):
+            continue
+        kids = list(c.named_children())
+        if not kids:
+            continue
+        has_list = any(isinstance(m, (nn.ModuleList, nn.Sequential)) for m in c.modules())
+        if has_list or len(kids) >= 2:
+            stages.append((n, c))
+    return stages
+
+
+def _capture_flow(model_obj, config, model_type) -> tuple[str, list[dict], list[int] | None] | None:
+    """Capture data-flow shapes via a forward on the ALREADY-BUILT *meta* model — meta tensors
+    carry shapes but allocate ~no real memory, so this is safe to run for every model in parallel.
+    Models whose forward can't run on meta (data-dependent ops) simply don't get the flow view;
+    they fall back to the structural module-tree view. We never build a real model here."""
+    try:
+        f = _run_flow(model_obj, config, "meta")
+        if f and len(f[1]) >= 2:
+            return f
+    except BaseException:
+        pass
+    return None
+
+
+_HEAD_PRIORITY = [
+    "IMAGE_TEXT_TO_TEXT",
+    "SEQ_TO_SEQ_CAUSAL_LM",
+    "CAUSAL_LM",
+    "IMAGE_CLASSIFICATION",
+    "AUDIO_CLASSIFICATION",
+    "MASKED_LM",
+    "SEQUENCE_CLASSIFICATION",
+    "TOKEN_CLASSIFICATION",
+    "PRETRAINING",
+]
+
+
+def _detect_family(mt: str, config):
+    """Return (family, head_class) by scanning the Auto MODEL_FOR_* mappings."""
+    try:
+        import transformers.models.auto.modeling_auto as ma
+    except Exception:
+        return None, None
+    for key in _HEAD_PRIORITY:
+        m = getattr(ma, f"MODEL_FOR_{key}_MAPPING_NAMES", {})
+        if mt in m:
+            hc = m[mt]
+            if isinstance(hc, (tuple, list)):  # some types map to several heads; take the first
+                hc = hc[0] if hc else None
+            return key.lower(), hc
+    return None, None
+
+
+def _attn_child(layer):
+    for n, c in layer.named_children():
+        if "attention" in type(c).__name__.lower() or n in ("self_attn", "attn", "attention", "self_attention"):
+            return c
+    return None
+
+
+def _mixer_child(layer):
+    """The token-mixing submodule of a non-attention block (Mamba/SSM/RWKV/linear-attention)."""
+    # prefer the conventionally-named mixer child
+    for n, c in layer.named_children():
+        if n in ("mixer", "ssm", "temporal_mixer", "time_mixer"):
+            return c
+    # else a non-norm child whose class name names a known mixer family
+    for n, c in layer.named_children():
+        cls = type(c).__name__
+        if _NORM_RE.search(cls):
+            continue
+        low = cls.lower()
+        if any(k in low for k in ("mamba", "mixer", "ssm", "rwkv", "retention", "lineardelta")):
+            return c
+    return None
+
+
+def _block_core(layer):
+    """The token-mixing module of a block: real attention if present, else an SSM/recurrent mixer."""
+    return _attn_child(layer) or _mixer_child(layer)
+
+
+def _annotate_moe(nodes: list[dict], am, ed) -> None:
+    """Fill experts/router counts (fused experts have no per-expert submodules on meta)."""
+    for nd in nodes:
+        if nd["kind"] == "experts" and not nd.get("dim"):
+            nd["dim"] = f"×{am.num_experts or '?'} · dim {ed or '?'}"
+        if nd["kind"] == "router":
+            nd["dim"] = f"top-{am.experts_per_token or '?'} of {am.num_experts or '?'}"
+
+
+def _mlp_child(layer):
+    for n, c in layer.named_children():
+        cl = type(c).__name__.lower()
+        if n in ("mlp", "feed_forward", "ffn", "block_sparse_moe", "moe") or cl.endswith(
+            ("mlp", "moe", "moeblock", "sparsemoeblock", "feedforward")
+        ):
+            return c
+    return None
+
+
+def _linear_proj(name, module) -> Proj | None:
+    import torch.nn as nn
+
+    if isinstance(module, nn.Linear):
+        return Proj(name=name, in_f=module.in_features, out_f=module.out_features, bias=module.bias is not None)
+    # some models use Conv1D (gpt2) with weight shape (in, out)
+    w = getattr(module, "weight", None)
+    if w is not None and hasattr(w, "shape") and len(w.shape) == 2:
+        return Proj(
+            name=name, in_f=int(w.shape[1]), out_f=int(w.shape[0]), bias=getattr(module, "bias", None) is not None
+        )
+    return None
+
+
+def _extract_block(layer, tcfg, am) -> BlockSpec | None:
+    """Decompose one decoder layer into a canonical pre-norm block for detailed rendering."""
+
+    spec = BlockSpec(layer_class=type(layer).__name__)
+    norms: list[tuple[str, str]] = []  # (child_name, class_name)
+    attn_mod = mlp_mod = None
+    for cname, child in layer.named_children():
+        cls = type(child).__name__
+        low = (cname + " " + cls).lower()
+        if _NORM_RE.search(cls):
+            norms.append((cname, cls))
+        elif low.find("attention") >= 0 or cname in ("self_attn", "attn", "attention", "self_attention"):
+            attn_mod = (cname, child)
+        elif cname in ("mlp", "feed_forward", "ffn", "block_sparse_moe", "moe", "mlp_layer") or cls.lower().endswith(
+            ("mlp", "moe", "moeblock", "sparsemoeblock", "feedforward")
+        ):
+            mlp_mod = (cname, child)
+
+    # ---- attention ----
+    if attn_mod is not None:
+        _, am_mod = attn_mod
+        aspec = AttentionSpec(cls=type(am_mod).__name__, variant=am.attn_variant)
+        aspec.n_heads = am.num_attention_heads
+        aspec.n_kv = am.num_kv_heads
+        aspec.head_dim = am.head_dim
+        aspec.sliding_window = am.sliding_window if isinstance(am.sliding_window, int) else None
+        for n, c in am_mod.named_children():
+            p = _linear_proj(n, c)
+            if p is not None:  # every Linear/Conv1D in attention is a projection (q/k/v/o, fused qkv, lora)
+                aspec.projs.append(p)
+            if _NORM_RE.search(type(c).__name__) and ("q_norm" in n or "k_norm" in n or n.endswith("norm")):
+                aspec.qk_norm = True
+            if "rotary" in type(c).__name__.lower():
+                aspec.rope = True
+        aspec.rope = aspec.rope or (am.positional == "RoPE")
+        spec.attention = aspec
+
+    # ---- mlp / moe ----
+    if mlp_mod is not None:
+        mname, mm = mlp_mod
+        mcls = type(mm).__name__
+        # MoE if this module (or any descendant) is a router/experts/sparse block
+        is_moe = bool(_MOE_RE.search(mcls)) or any(_MOE_RE.search(type(c).__name__) for c in mm.modules())
+        mspec = MLPSpec(cls=mcls, is_moe=is_moe, act=(am.hidden_act or "SiLU"))
+        if is_moe:
+            mspec.n_experts = am.num_experts
+            mspec.top_k = am.experts_per_token
+            mspec.expert_dim = _cfg(tcfg, "moe_intermediate_size", "expert_intermediate_size", "intermediate_size")
+            mspec.n_shared = _cfg(tcfg, "n_shared_experts", "num_shared_experts")  # counts only
+        else:
+            projs = {}
+            for n, c in mm.named_children():
+                p = _linear_proj(n, c)
+                if p is not None:
+                    projs[n.lower()] = p
+            mspec.projs = list(projs.values())
+            mspec.gate = next((projs[k] for k in projs if "gate" in k and "up" not in k), None)
+            mspec.up = next(
+                (projs[k] for k in projs if k.startswith("up") or "up_proj" in k or "fc" in k or "c_fc" in k), None
+            )
+            mspec.down = next(
+                (
+                    projs[k]
+                    for k in projs
+                    if "down" in k
+                    or "c_proj" in k
+                    or k.endswith(("proj2", "_proj"))
+                    and "up" not in k
+                    and "gate" not in k
+                ),
+                None,
+            )
+        spec.mlp = mspec
+
+    # ---- norms (canonical pre-norm: input -> attn, post-attn -> mlp) ----
+    if norms:
+        pre = next((c for n, c in norms if "input" in n or "pre" in n or n == norms[0][0]), norms[0][1])
+        spec.pre_attn_norm = pre
+        post = next((c for n, c in norms if "post" in n or "attention" in n), None)
+        spec.post_attn_norm = post or (norms[1][1] if len(norms) > 1 else pre)
+        if len(norms) > 2:
+            spec.extra.append(f"{len(norms)} norms (sandwich/extra)")
+
+    # ---- token mixer (Mamba/SSM/RWKV) when there is no attention ----
+    if spec.attention is None:
+        mix = _mixer_child(layer)
+        if mix is not None:
+            mspec2 = AttentionSpec(cls=type(mix).__name__, variant="SSM")
+            for n, c in mix.named_children():
+                p = _linear_proj(n, c)
+                if p is not None:
+                    mspec2.projs.append(p)
+            spec.mixer = mspec2
+
+    if spec.attention is None and spec.mixer is None and spec.mlp is None:
+        return None
+    return spec
+
+
+_EXAMPLE_PROMPT = "Hey, how are you?"
+_FALLBACK_TOKENS = ["Hey", ",", "␣how", "␣are", "␣you", "?"]
+
+
+def _example_tokens(checkpoint: str | None) -> list[str]:
+    """Tokenize a fixed example prompt with the model's tokenizer (offline / cached only).
+
+    Falls back to a fixed illustrative token list when no cached tokenizer is available.
+    Token strings are cleaned: the SentencePiece ``▁`` and BPE ``Ġ`` space markers become a
+    visible ``␣`` so the per-token boxes read naturally.
+    """
+    if not checkpoint:
+        return list(_FALLBACK_TOKENS)
+    try:
+        from transformers import AutoTokenizer
+
+        tok = AutoTokenizer.from_pretrained(checkpoint, local_files_only=True)
+        ids = tok(_EXAMPLE_PROMPT, add_special_tokens=False)["input_ids"]
+        pieces = tok.convert_ids_to_tokens(ids)
+        out = []
+        for p in pieces[:10]:
+            out.append(p.replace("▁", "␣").replace("Ġ", "␣").replace("Ċ", "⏎"))
+        return out or list(_FALLBACK_TOKENS)
+    except Exception:
+        return list(_FALLBACK_TOKENS)
+
+
+_AUTO_RE = re.compile(r"\b(AutoModel(?:ForCausalLM|ForConditionalGeneration|ForImageTextToText)?|AutoModelFor\w+)\b")
+
+
+def _auto_classes(model: str) -> list[str]:
+    """Which Auto* classes the modeling file uses to instantiate sub-models (offline grep).
+
+    Composition models build their towers with e.g. ``AutoModel.from_config(vision_config)``;
+    surfacing this is important — the vision tower or LLM is often an AutoModel.
+    """
+    import os
+
+    from .discover import models_root
+
+    d = os.path.join(models_root(), model)
+    if not os.path.isdir(d):
+        return []
+    found: set[str] = set()
+    for f in sorted(os.listdir(d)):
+        if f.startswith(("modeling_", "modular_")) and f.endswith(".py"):
+            try:
+                src = open(os.path.join(d, f), encoding="utf-8").read()
+            except Exception:
+                continue
+            for m in _AUTO_RE.findall(src):
+                found.add(m)
+    return sorted(found)
+
+
+_VISION_RE = re.compile(r"(vision|visual|image|vit|clip|siglip|patch)", re.IGNORECASE)
+_AUDIO_RE = re.compile(r"(audio|speech|wav|mel|acoustic|whisper)", re.IGNORECASE)
+_PROJ_RE = re.compile(r"(projector|connector|multi_modal|mm_proj|adapter|merger|aligner)", re.IGNORECASE)
+_LLM_RE = re.compile(r"(language_model|text_model|^model$|decoder|thinker)", re.IGNORECASE)
+
+
+def _subconfig(config, *names):
+    for n in names:
+        sc = getattr(config, n, None)
+        if sc is not None and hasattr(sc, "model_type"):
+            return sc
+    return None
+
+
+def _tower_summary(child_name, child, config) -> dict:
+    """Summarize one top-level modality tower (encoder/projector) for the multi-tower view."""
+    cls = type(child).__name__
+    role = "other"
+    nl = child_name.lower() + " " + cls.lower()
+    if _PROJ_RE.search(nl):
+        role = "projector"
+    elif _VISION_RE.search(nl):
+        role = "vision"
+    elif _AUDIO_RE.search(nl):
+        role = "audio"
+    elif _LLM_RE.search(child_name.lower()):
+        role = "text"
+    sc = None
+    if role == "vision":
+        sc = _subconfig(config, "vision_config")
+    elif role == "audio":
+        sc = _subconfig(config, "audio_config")
+    layers = hidden = mtype = None
+    if sc is not None:
+        layers = _cfg(sc, "num_hidden_layers", "depth", "num_layers")
+        hidden = _cfg(sc, "hidden_size", "d_model", "embed_dim")
+        mtype = getattr(sc, "model_type", None)
+    out = {"role": role, "name": child_name, "cls": cls, "model_type": mtype, "layers": layers, "hidden": hidden}
+    # decompose small towers (projector / connector) fully so their inside is visible.
+    n_sub = sum(1 for _ in child.modules())
+    if role == "projector" or (role in ("vision", "audio") and n_sub <= 12):
+        out["children"] = _decompose(child, 2)
+    elif role in ("vision", "audio"):
+        # big encoder stack: show ONE representative layer (decomposed) ×N, like the LLM block
+        import contextlib
+
+        with contextlib.suppress(Exception):
+            _, elayers = _find_layer_list(child, layers)
+            if elayers is not None and len(elayers) > 0:
+                rep = elayers[0]
+                if _attn_child(rep) is None:  # hierarchical (e.g. swin stage) -> descend
+                    rep = next((m for m in rep.modules() if _attn_child(m) is not None), rep)
+                out["block_children"] = _decompose(rep, 2)
+                out["block_class"] = type(rep).__name__
+                out["block_n"] = len(elayers)
+    return out
+
+
+def _detect_towers(model_obj, config) -> list[dict]:
+    """Detect modality towers from the top-level children of a (multimodal) model."""
+    towers = []
+    for n, c in model_obj.named_children():
+        cls = type(c).__name__
+        # skip plain containers like embeddings/heads at the very top
+        if cls in ("Embedding", "Linear") and not _PROJ_RE.search(n.lower()):
+            continue
+        t = _tower_summary(n, c, config)
+        if t["role"] in ("vision", "audio", "projector", "text"):
+            towers.append(t)
+    return towers
+
+
+def introspect(model: str, model_type: str | None = None) -> ArchModel:
+    """Introspect ``model`` into an ``ArchModel``. Never raises -- failures are recorded."""
+    from transformers.models.auto.configuration_auto import CONFIG_MAPPING
+
+    am = ArchModel(model=model, model_type=model_type)
+    mt = model_type or model
+    try:
+        config_cls = CONFIG_MAPPING[mt]
+    except Exception:
+        # try to recover model_type from the directory name
+        config_cls = None
+        for k, cls in CONFIG_MAPPING.items():
+            if k.replace("-", "_") == model:
+                config_cls, mt = cls, k
+                break
+        if config_cls is None:
+            am.status = "failed"
+            am.error = f"no config class for model_type={mt!r}"
+            return am
+
+    am.model_type = mt
+    am.config_class = config_cls.__name__
+    try:
+        config = config_cls()
+    except Exception as e:  # some configs require args
+        am.status = "failed"
+        am.error = f"config instantiation failed: {e!r}"[:300]
+        return am
+
+    # decoder/text config (multimodal configs nest a text_config)
+    tcfg = getattr(config, "get_text_config", lambda: config)()
+    am.cross_attention_layers = _cfg(tcfg, "cross_attention_layers")
+    has_vision_cfg = _subconfig(config, "vision_config") is not None
+    has_audio_cfg = _subconfig(config, "audio_config") is not None
+
+    # task family + which layout view to use (decoder LM / encoder-classifier / enc-dec / VLM)
+    am.family, am.head_class = _detect_family(mt, config)
+    am.num_labels = _cfg(config, "num_labels", "num_classes")
+    if getattr(config, "is_encoder_decoder", False):
+        am.view = "enc_dec"
+    elif am.family in ("image_classification", "audio_classification", "masked_lm"):
+        am.view = "encoder"
+    else:
+        am.view = "decoder"  # multimodal is set later once towers are detected
+
+    am.hidden_size = _cfg(tcfg, "hidden_size", "d_model", "n_embd", "dim")
+    am.vocab_size = _cfg(tcfg, "vocab_size")
+    am.num_layers = _cfg(tcfg, "num_hidden_layers", "num_layers", "n_layer", "n_layers")
+    am.num_attention_heads = _cfg(tcfg, "num_attention_heads", "n_head", "num_heads")
+    am.num_kv_heads = _cfg(tcfg, "num_key_value_heads", "num_kv_heads", default=am.num_attention_heads)
+    am.head_dim = _cfg(tcfg, "head_dim")
+    # hierarchical vision models (swin, dinat, hiera, ...) carry per-stage tuples; only the
+    # plain int case has a well-defined single head_dim.
+    if am.head_dim is None and isinstance(am.hidden_size, int) and isinstance(am.num_attention_heads, int):
+        am.head_dim = am.hidden_size // am.num_attention_heads
+    am.intermediate_size = _cfg(tcfg, "intermediate_size", "ffn_dim", "d_ff")
+    am.max_position_embeddings = _cfg(tcfg, "max_position_embeddings", "n_positions")
+    am.sliding_window = _cfg(tcfg, "sliding_window")
+    am.num_experts = _cfg(tcfg, "num_experts", "num_local_experts", "n_routed_experts", "moe_num_experts")
+    am.experts_per_token = _cfg(tcfg, "num_experts_per_tok", "moe_topk", "num_experts_per_token", "top_k")
+    am.tie_word_embeddings = _cfg(config, "tie_word_embeddings")
+    am.attn_variant = _attn_variant(tcfg)
+    am.hidden_act = _cfg(tcfg, "hidden_act", "activation_function", "hidden_activation")
+    rt = _cfg(tcfg, "rope_theta")
+    am.rope_theta = rt if isinstance(rt, (int, float)) else None
+    lts = _cfg(tcfg, "layer_types")
+    if isinstance(lts, (list, tuple)) and lts and all(isinstance(x, str) for x in lts):
+        am.layer_types = list(lts)
+    am.checkpoint = _example_checkpoint(model, mt)
+    am.tokens = _example_tokens(am.checkpoint)
+    # sparse/compressed attention components, keyed by distinct attention layer_type
+    compress_rates = _cfg(tcfg, "compress_rates")
+    if isinstance(compress_rates, dict) and am.layer_types:
+        index_topk = _cfg(tcfg, "index_topk")
+        from .masks import concat_compressed_mask
+
+        for lt in dict.fromkeys(am.layer_types):
+            if lt in compress_rates:
+                comp = {"type": lt, "compress_rate": compress_rates[lt]}
+                if "sparse" in lt and index_topk:
+                    comp["index_topk"] = index_topk
+                # the real mask passed to attention: sliding K/V cache ⊕ compressed cache.
+                # display m is capped so the figure is legible (true rate kept in the label).
+                m = min(int(compress_rates[lt]), 8)
+                grid, split, n_comp = concat_compressed_mask(16, m)
+                comp.update(display_m=m, mask=grid, mask_split=split, n_comp=n_comp)
+                am.sparse_components.append(comp)
+
+    module_classes: set[str] = set()
+    layer_classes: list[str] = []
+
+    try:
+        model_obj = _build_meta_model(config)
+        am.top_class = type(model_obj).__name__
+        counts: dict[str, int] = {}
+        for _, mod in model_obj.named_modules():
+            cn = type(mod).__name__
+            counts[cn] = counts.get(cn, 0) + 1
+            module_classes.add(cn)
+        am.module_kinds = dict(sorted(counts.items(), key=lambda kv: -kv[1])[:25])
+
+        name, layers = _find_layer_list(model_obj, am.num_layers)
+
+        # multimodal towers (vision/audio encoders + projector + LLM)
+        towers = _detect_towers(model_obj, config)
+        if (has_vision_cfg or has_audio_cfg) and any(t["role"] in ("vision", "audio") for t in towers):
+            am.is_multimodal = True
+            am.view = "multimodal"
+            am.towers = towers
+            am.modal_inputs = ["input_ids"]
+            if any(t["role"] == "vision" for t in towers):
+                am.modal_inputs.append("pixel_values")
+            if any(t["role"] == "audio" for t in towers):
+                am.modal_inputs.append("input_features")
+            am.auto_classes = _auto_classes(model)
+            # prefix-LM image attention (bidirectional over image tokens, e.g. PaliGemma/Gemma3)
+            import os as _os
+
+            from .discover import models_root as _mr
+
+            _d = _os.path.join(_mr(), model)
+            for _f in sorted(_os.listdir(_d)) if _os.path.isdir(_d) else []:
+                if _f.startswith(("modeling_", "modular_")) and _f.endswith(".py"):
+                    try:
+                        _src = open(_os.path.join(_d, _f), encoding="utf-8").read()
+                    except Exception:
+                        continue
+                    if "token_type_ids" in _src or "bidirectional" in _src.lower():
+                        am.image_bidirectional = True
+                        break
+            # for the LLM block, restrict the layer search to the language_model subtree
+            llm = next((getattr(model_obj, t["name"]) for t in towers if t["role"] == "text"), None)
+            if llm is not None:
+                _, llm_layers = _find_layer_list(llm, am.num_layers)
+                if llm_layers is not None:
+                    layers = llm_layers
+
+        # norm/positional must be known before block extraction (rope detection uses them)
+        am.norm_type = _norm_type(module_classes)
+        am.positional = _positional(tcfg, module_classes)
+        if layers is not None:
+            layer_classes = [type(layer).__name__ for layer in layers]
+            if not am.num_layers:
+                am.num_layers = len(layers)
+            # decompose a representative decoder layer into block internals (full view).
+            # MoE models often keep the first few layers dense (e.g. DeepSeek); prefer a
+            # layer that actually contains the MoE block so the diagram is representative.
+            try:
+                import torch.nn as nn
+
+                def _score(layer):
+                    has_attn = any(
+                        "attention" in type(c).__name__.lower() and any(isinstance(g, nn.Linear) for g in c.children())
+                        for c in layer.children()
+                    )
+                    has_moe = any(_MOE_RE.search(type(c).__name__) for c in layer.modules())
+                    # prefer a layer with real attention; among those prefer one with MoE
+                    return (1 if has_attn else 0, 1 if (has_moe and am.num_experts) else 0)
+
+                rep = layers[0]
+                if len(layers) > 1:
+                    rep = max(layers, key=_score)
+                # hierarchical/staged models (Swin/Donut/NAT): the "layer" is a Stage with no
+                # direct attention -> descend to the innermost module that has real attention
+                if _attn_child(rep) is None:
+                    deep = next(
+                        (m for m in rep.modules() if _attn_child(m) is not None and m is not rep),
+                        None,
+                    )
+                    if deep is None:  # search the whole model
+                        deep = next((m for m in model_obj.modules() if _attn_child(m) is not None), None)
+                    if deep is not None:
+                        rep = deep
+                        layers = [rep]
+                am.block = _extract_block(rep, tcfg, am)
+                # record any structurally different layer classes (hybrid) not shown in detail
+                kinds = {_kind_for_layer(type(layer).__name__) for layer in layers}
+                am.alt_blocks = sorted(k for k in kinds if k != "attention")
+
+                # FULL recursive decomposition of each DISTINCT attention variant (the real
+                # module tree -- e.g. DeepSeek-V4 HCA vs CSA-with-Indexer), so all blocks show.
+                variants, seen = [], set()
+                for i, layer in enumerate(layers):
+                    at = _block_core(layer)
+                    if at is None:
+                        continue
+                    lt = getattr(at, "layer_type", None) or (
+                        am.layer_types[i] if i < len(am.layer_types) else type(at).__name__
+                    )
+                    if lt in seen:
+                        continue
+                    seen.add(lt)
+                    variants.append({"type": lt, "cls": type(at).__name__, "children": _decompose(at, 2)})
+                    if len(variants) >= 3:
+                        break
+                am.attention_variants = variants
+                # decompose the FFN/MoE of the representative layer too
+                for n, c in rep.named_children():
+                    if n in ("mlp", "feed_forward", "ffn", "block_sparse_moe", "moe") or type(
+                        c
+                    ).__name__.lower().endswith(("mlp", "moe", "moeblock", "sparsemoeblock")):
+                        am.mlp_tree = {"cls": type(c).__name__, "children": _decompose(c, 2)}
+                        # annotate router/experts nodes with counts (fused experts have no
+                        # per-expert submodules on meta, so surface the count + dims here)
+                        ed = _cfg(tcfg, "moe_intermediate_size", "expert_intermediate_size", "intermediate_size")
+                        _annotate_moe(am.mlp_tree["children"], am, ed)
+                        break
+
+                # distinct FFN variants per config.mlp_layer_types (moe / hash_moe / dense)
+                mlts = _cfg(tcfg, "mlp_layer_types")
+                if isinstance(mlts, (list, tuple)) and mlts and all(isinstance(x, str) for x in mlts):
+                    am.mlp_layer_types = list(mlts)
+                    ed = _cfg(tcfg, "moe_intermediate_size", "expert_intermediate_size", "intermediate_size")
+                    mseen = set()
+                    for i, layer in enumerate(layers):
+                        if i >= len(am.mlp_layer_types):
+                            break
+                        mt2 = am.mlp_layer_types[i]
+                        if mt2 in mseen:
+                            continue
+                        mc = _mlp_child(layer)
+                        if mc is None:
+                            continue
+                        mseen.add(mt2)
+                        ch = _decompose(mc, 2)
+                        _annotate_moe(ch, am, ed)
+                        am.mlp_variants.append({"type": mt2, "cls": type(mc).__name__, "children": ch})
+                        if len(am.mlp_variants) >= 3:
+                            break
+            except Exception:
+                am.block = None
+
+        # Universal fallback: any model without a recognized attention/SSM/MLP transformer block
+        # (conv backbones, audio codecs, FFT mixers, detectors, ...) still gets a real top-down
+        # render from its actual (collapsed) module tree instead of a config-only schematic.
+        no_block = am.block is None or (am.block.attention is None and am.block.mixer is None and am.block.mlp is None)
+        generic = no_block and not am.attention_variants
+        try:
+            am.is_pipeline = len(_stage_children(model_obj)) >= 3
+        except Exception:
+            am.is_pipeline = False
+        # A model gets the data-flow view when it has no single transformer block (conv backbone,
+        # codec, ...) OR when it is a multi-stage pipeline (BLT: patcher → encoder → transformer →
+        # decoder), where the data path is the story rather than one attention block.
+        if (generic or am.is_pipeline) and not am.is_multimodal:
+            try:
+                am.generic_tree = am.generic_tree or _module_tree(model_obj, depth=4)
+            except Exception:
+                am.generic_tree = []
+            try:
+                cap = _capture_flow(model_obj, config, model)
+                if cap:
+                    am.flow_input, am.flow, am.flow_output = cap
+            except Exception:
+                am.flow = []
+    except Exception as e:
+        am.status = "config-only"
+        am.error = f"meta build failed: {e!r}"[:300]
+
+    # MoE / decoder kind
+    moe_modules = {c for c in module_classes if _MOE_RE.search(c)}
+    am.is_moe = bool(moe_modules) or (am.num_experts is not None and am.num_experts > 1)
+    has_mamba = any(_MAMBA_RE.search(c) for c in module_classes)
+    has_linear = any(_LINEAR_ATTN_RE.search(c) for c in module_classes)
+
+    if am.norm_type is None:
+        am.norm_type = _norm_type(module_classes)
+    if am.positional is None:
+        am.positional = _positional(tcfg, module_classes)
+
+    # layer blocks: prefer config.layer_types (the real per-layer schedule) when present,
+    # since hybrid models (qwen3_next: linear+full, gemma3: sliding+full) reuse ONE layer
+    # class so the module tree alone hides the mix.
+    blocks: list[LayerBlock] = []
+    if am.layer_types:
+        for lt, n in _collapse(am.layer_types):
+            kind = layer_type_kind(lt)
+            blocks.append(
+                LayerBlock(
+                    layer_class=lt,
+                    count=n,
+                    kind=kind,
+                    attn_variant=(am.attn_variant if kind == "attention" else None),
+                    mlp_kind="moe" if (am.is_moe and kind in ("attention", "linear_attention")) else "dense",
+                    norm=am.norm_type,
+                )
+            )
+    elif layer_classes:
+        for cls_name, n in _collapse(layer_classes):
+            kind = _kind_for_layer(cls_name)
+            blocks.append(
+                LayerBlock(
+                    layer_class=cls_name,
+                    count=n,
+                    kind=kind,
+                    attn_variant=am.attn_variant if kind == "attention" else None,
+                    mlp_kind="moe" if (am.is_moe and kind == "attention") else "dense",
+                    norm=am.norm_type,
+                )
+            )
+    elif am.num_layers:
+        # config-only fallback: one synthetic block. With no module tree to walk, infer the
+        # kind from the synthetic layer name (e.g. "MambaDecoderLayer" -> mamba) and config.
+        synth_name = f"{am.config_class.replace('Config', '') if am.config_class else mt}DecoderLayer"
+        kind = "mamba" if has_mamba else ("linear_attention" if has_linear else _kind_for_layer(synth_name))
+        blocks.append(
+            LayerBlock(
+                layer_class=synth_name,
+                count=am.num_layers,
+                kind=kind,
+                attn_variant=am.attn_variant,
+                mlp_kind="moe" if am.is_moe else "dense",
+                norm=am.norm_type,
+            )
+        )
+    am.layer_blocks = blocks
+
+    # decoder kind + summary
+    distinct_kinds = {b.kind for b in blocks}
+    if len(distinct_kinds) > 1:
+        am.decoder_kind = "hybrid"
+    elif has_mamba and not any(b.kind == "attention" for b in blocks):
+        am.decoder_kind = "ssm"
+    elif am.is_moe:
+        am.decoder_kind = "sparse-moe"
+    elif distinct_kinds and distinct_kinds != {"attention"}:
+        am.decoder_kind = "other"
+    else:
+        am.decoder_kind = "dense"
+
+    if blocks:
+        # aggregate by kind-label totals (not per-run) so the summary stays short
+        totals: dict[str, int] = {}
+        for b in blocks:
+            label = _pretty_kind(b)
+            totals[label] = totals.get(label, 0) + b.count
+        am.layer_summary = " + ".join(f"{n}× {label}" for label, n in totals.items())
+
+    # attention-mask patterns per distinct layer type (offline, real masking_utils)
+    if am.layer_types:
+        try:
+            from .masks import attention_pattern_grids
+
+            am.attn_patterns = attention_pattern_grids(tcfg, am.layer_types, seq=24)
+        except Exception:
+            am.attn_patterns = {}
+
+    return am
+
+
+_LT_SHORT = {
+    "full_attention": "full",
+    "sliding_attention": "sliding",
+    "chunked_attention": "chunked",
+    "compressed_sparse_attention": "compressed-sparse",
+    "heavily_compressed_attention": "heavily-compressed",
+    "linear_attention": "linear",
+}
+
+
+def _short_layer_type(name: str) -> str:
+    if name in _LT_SHORT:
+        return _LT_SHORT[name]
+    n = name.lower()
+    for key, short in _LT_SHORT.items():
+        if n == key:
+            return short
+    return name.replace("_attention", "").replace("_", " ")
+
+
+def _pretty_kind(b: LayerBlock) -> str:
+    # when the layer_class is actually a layer_type (sliding/full/...), show that label
+    is_layer_type = b.layer_class.endswith("attention") or b.layer_class in _LT_SHORT
+    if b.kind == "attention":
+        base = _short_layer_type(b.layer_class) if is_layer_type else (b.attn_variant or "Attention")
+        return base + (" +MoE" if b.mlp_kind == "moe" else "")
+    if b.kind == "linear_attention":
+        base = _short_layer_type(b.layer_class) if is_layer_type else "linear-attn"
+        return base + (" +MoE" if b.mlp_kind == "moe" else "")
+    return {"mamba": "Mamba", "recurrent": "Recurrent"}.get(b.kind, b.layer_class)

arch_svg/layout.py

@@ -0,0 +1,1479 @@
+"""Turn an ``ArchModel`` / ``ArchDiff`` into positioned boxes.
+
+Layout is deterministic (stable ordering, integer coordinates) so the emitted SVG is
+diff-able in git. The visual vocabulary is *standardized*: a given component kind always
+maps to the same shape class and palette slot, mirroring the library's "standardize, don't
+abstract" tenet -- an attention block looks the same across every model's diagram.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from .engine import text_width
+from .introspect import ArchModel
+from .modular import ArchDiff, ClassChange
+
+
+_PARENT_CACHE: dict = {}  # model_type -> ArchModel of the modular parent (for diff comparison)
+
+
+def _parent_arch(model_type):
+    if model_type not in _PARENT_CACHE:
+        try:
+            from .introspect import introspect
+
+            _PARENT_CACHE[model_type] = introspect(model_type)
+        except Exception:
+            _PARENT_CACHE[model_type] = None
+    return _PARENT_CACHE[model_type]
+
+
+def _collect_submodule_names(nodes, out: set):
+    for nd in nodes:
+        out.add(nd["name"])
+        if nd.get("children"):
+            _collect_submodule_names(nd["children"], out)
+
+
+# component kind -> css class used by render.py (the standardized vocabulary)
+KIND_CLASS = {
+    "embedding": "c-embed",
+    "attention": "c-attn",
+    "mamba": "c-mamba",
+    "linear_attention": "c-linattn",
+    "recurrent": "c-recur",
+    "moe": "c-moe",
+    "mlp": "c-mlp",
+    "norm": "c-norm",
+    "head": "c-head",
+    "config": "c-config",
+    "rope": "c-rope",
+    "layer": "c-layer",
+}
+
+# diff change -> css class
+CHANGE_CLASS = {"added": "ch-added", "overridden": "ch-over", "deleted": "ch-deleted", None: ""}
+
+
+@dataclass
+class Box:
+    x: int
+    y: int
+    w: int
+    h: int
+    label: str
+    cls: str = ""
+    sublabels: list[str] = field(default_factory=list)
+    badge: str | None = None  # small corner tag, e.g. "×32"
+    ghost: bool = False
+    change: str | None = None  # added | overridden | deleted | None
+    title: str | None = None  # SVG <title> tooltip
+    shape: str = "rect"  # rect | circle | container | op | io | cell | grid
+    small: bool = False  # render with smaller fonts (inner ops, Q/K/V chips)
+    glyph: str | None = None  # single char drawn centered (e.g. "+", "×")
+    grid: list[list[int]] | None = None  # for shape == "grid": 0/1 attention-mask matrix
+    grid_split: int | None = None  # column index dividing sliding K/V | compressed cache
+
+
+@dataclass
+class Arrow:
+    """A poly-line connector. ``points`` are (x, y) pairs; rendered with an arrowhead."""
+
+    points: list[tuple[int, int]]
+    cls: str = "flow"  # flow | residual
+    label: str | None = None
+    dashed: bool = False
+
+
+@dataclass
+class Diagram:
+    width: int
+    height: int
+    boxes: list[Box]
+    title: str
+    subtitle: str = ""
+    legend: list[tuple[str, str]] = field(default_factory=list)  # (css-class, text)
+    facts: list[tuple[str, str]] = field(default_factory=list)  # side panel key/value
+    changes: list[tuple[str, str, str]] = field(default_factory=list)  # (type, class, detail)
+    arrows: list[Arrow] = field(default_factory=list)
+    mode: str = "full"
+    spine: bool = True  # draw the central connector spine
+
+
+# geometry constants
+W = 920
+COL_X = 150  # left edge of the central stack
+COL_W = 420  # width of the central stack
+PAD = 28
+ROW_H = 30
+
+
+def _fmt_val(v) -> str:
+    """Compactly format a config value -- including sequence-valued fields.
+
+    Some (often hybrid) models carry per-layer lists for ``intermediate_size``,
+    ``head_dim`` or ``sliding_window``. We collapse those to ``[x]×N`` (all equal) or
+    ``[a … z] (N)`` rather than dumping the whole list. No per-model handling.
+    """
+    if v is None:
+        return "—"
+    if isinstance(v, (list, tuple)):
+        if not v:
+            return "[]"
+        if len(set(v)) == 1:
+            return f"[{_fmt_val(v[0])}]×{len(v)}"
+        return f"[{_fmt_val(v[0])} … {_fmt_val(v[-1])}] ({len(v)})"
+    if isinstance(v, bool):
+        return str(v)
+    if isinstance(v, int):
+        return f"{v:,}"
+    return str(v)
+
+
+def _facts(am: ArchModel) -> list[tuple[str, str]]:
+    fmt = _fmt_val
+
+    mid = am.checkpoint or am.model
+    if len(mid) > 30:
+        mid = mid[:29] + "…"
+    rows = [
+        ("model id", mid),
+        ("model_type", am.model_type or "—"),
+        ("config", am.config_class or "—"),
+        ("decoder", am.decoder_kind),
+        ("layers", fmt(am.num_layers)),
+        ("hidden", fmt(am.hidden_size)),
+        ("intermediate", fmt(am.intermediate_size)),
+        ("heads / kv", f"{am.num_attention_heads} / {am.num_kv_heads}"),
+        ("attention", am.attn_variant or "—"),
+        ("norm", am.norm_type or "—"),
+        ("positional", am.positional or "—"),
+        ("vocab", fmt(am.vocab_size)),
+        ("max_pos", fmt(am.max_position_embeddings)),
+    ]
+    if am.sliding_window:
+        rows.append(("sliding_window", fmt(am.sliding_window)))
+    if am.is_moe:
+        rows.append(("experts (top-k)", f"{am.num_experts} (top-{am.experts_per_token})"))
+    rows.append(("tie_embeddings", str(am.tie_word_embeddings)))
+    if am.status != "ok":
+        rows.append(("status", am.status))
+    return rows
+
+
+def _reflow(boxes: list[Box], start_y: int = PAD + 70, gap: int = 16) -> int:
+    """Re-stack boxes vertically in list order using their (possibly grown) heights.
+
+    Called after any box height changes so diff annotations never overlap neighbours.
+    Returns the y coordinate just below the last box.
+    """
+    y = start_y
+    for b in boxes:
+        b.y = y
+        y += b.h + gap
+    return y
+
+
+def _build_compact(am: ArchModel) -> Diagram:
+    """Top-down view from CONFIG alone — used when the model can't be instantiated on meta (so
+    there's no module tree to decompose) or has no transformer block (e.g. conv backbones).
+    Shows input → embedding → the per-config layer blocks (collapsed, ×N) → final norm → head,
+    family-aware, with the layer-type schedule. Far cleaner than the old bottom-up stack."""
+    boxes: list[Box] = []
+    x, w = COL_X, COL_W
+    y = PAD + 78
+    modality = (
+        "vision"
+        if am.family == "image_classification"
+        else ("audio" if am.family == "audio_classification" else "text")
+    )
+    enc = am.view == "encoder"
+
+    def add(label, cls, subs=None, h=None, shape="op", badge=None):
+        nonlocal y
+        b = Box(
+            x,
+            y,
+            w,
+            h if h is not None else (30 + 15 * len(subs or [])),
+            label,
+            cls,
+            sublabels=subs or [],
+            shape=shape,
+            small=True,
+            badge=badge,
+        )
+        boxes.append(b)
+        y = y + b.h + 16
+        return b
+
+    # input + embedding (family-aware)
+    if modality == "vision":
+        add("input image", "c-io", ["pixel_values [1, 3, H, W]"], h=30, shape="io")
+        add("Patch Embedding", "c-embed", [f"→ [1, S, {am.hidden_size or '?'}]"], h=40)
+    elif modality == "audio":
+        add("input audio", "c-io", ["input_features"], h=30, shape="io")
+        add("Feature Projection", "c-embed", [f"→ [1, S, {am.hidden_size or '?'}]"], h=40)
+    else:
+        add("input_ids", "c-io", ["[1, S]"], h=30, shape="io")
+        add("Token Embedding", "c-embed", [f"[{am.vocab_size or '?'} × {am.hidden_size or '?'}]"], h=40)
+
+    # per-config layer blocks (collapsed ×N), coloured by kind
+    for b in am.layer_blocks:
+        inner = [f"{b.norm or 'Norm'} → {(b.attn_variant or b.kind.replace('_', '-'))}"]
+        inner.append(
+            f"{b.norm or 'Norm'} → {('MoE ' + str(am.num_experts) + 'E·top-' + str(am.experts_per_token)) if b.mlp_kind == 'moe' else 'MLP'}"
+        )
+        cls = KIND_CLASS["moe"] if b.mlp_kind == "moe" else KIND_CLASS.get(b.kind, KIND_CLASS["layer"])
+        add(b.layer_class, cls, inner, badge=f"×{b.count}", h=30 + 15 * len(inner))
+
+    add(f"Final {am.norm_type or 'Norm'}", "c-norm", h=28)
+    if enc:
+        add("Pool + Classifier head", "c-head", [f"→ class logits [1, {am.num_labels or '?'}]"], h=40)
+    else:
+        add("LM Head", "c-head", [f"→ logits [1, S, {am.vocab_size or '?'}]"], h=40)
+
+    facts = _facts(am)
+    height = max(y + PAD, PAD + 78 + len(facts) * 22 + 120)
+    legend = _STD_LEGEND
+    note = (
+        "config-only — built from config (weights not instantiable on meta)"
+        if am.status != "ok"
+        else "no transformer block (conv backbone) — shown from config"
+    )
+    return Diagram(
+        width=W,
+        height=height,
+        boxes=boxes,
+        title=am.model,
+        subtitle=f"⚠ {note}  ·  {am.decoder_kind} · {am.layer_summary or ''}",
+        legend=legend,
+        facts=facts,
+        mode="full",
+        spine=True,
+    )
+
+
+def _build_generic(am: ArchModel) -> Diagram:
+    """Universal top-down render from the real (collapsed) module tree — for any model that is
+    not a standard attention/SSM transformer: conv backbones (ResNet/ConvNeXt/MobileNet),
+    audio codecs (EnCodec/DAC), FFT mixers (FNet), detectors, pose/segmentation heads, ...
+    Every ``ModuleList``/``Sequential`` of identical blocks is collapsed to a ``×N`` section."""
+    boxes: list[Box] = []
+    GX, GW = 40, 850  # tree column (facts panel sits at width-288)
+    y = PAD + 78
+    PFX = (am.config_class or "").replace("Config", "")
+
+    def add(label, cls, subs=None, h=30, shape="op"):
+        nonlocal y
+        b = Box(GX, y, GW, h, label, cls, sublabels=subs or [], shape=shape, small=True)
+        boxes.append(b)
+        y = b.y + b.h + 14
+        return b
+
+    # input (family-aware) — what the model consumes
+    modality = (
+        "vision" if am.family == "image_classification" else ("audio" if am.family == "audio_classification" else None)
+    )
+    if "pixel_values" in (am.modal_inputs or []) or modality == "vision":
+        add("input image", "c-io", ["pixel_values [1, 3, H, W]"], shape="io")
+    elif "input_features" in (am.modal_inputs or []) or modality == "audio":
+        add("input audio", "c-io", ["input_values / input_features"], shape="io")
+    else:
+        add("inputs", "c-io", ["input_ids / pixel_values / input_values"], shape="io")
+
+    # the real module hierarchy
+    y = _emit_tree(boxes, am.generic_tree, GX, GW, y, gp=8, prefix=PFX)
+
+    facts = _facts(am)
+    height = max(y + PAD, PAD + 78 + len(facts) * 22 + 120)
+    legend = _STD_LEGEND
+    return Diagram(
+        width=WF,
+        height=height,
+        boxes=boxes,
+        title=am.model,
+        subtitle=f"{am.top_class or ''}  ·  full module tree (no attention — shown structurally)",
+        legend=legend,
+        facts=facts,
+        mode="full",
+        spine=False,
+    )
+
+
+def _struct_hint(node: dict | None) -> str | None:
+    """A one-line structural summary of a stage's children, e.g. '4× DacEncoderBlock'."""
+    if not node:
+        return None
+    reps = [c["cls"] for c in node.get("children", []) if "×" in c.get("cls", "")]
+    if reps:
+        return " · ".join(reps[:2])
+    kids = node.get("children", [])
+    if kids:
+        return f"{len(kids)} submodules"
+    return None
+
+
+def _shape_str(s) -> str:
+    return "×".join(str(d) for d in s) if s else "?"
+
+
+def _build_flow(am: ArchModel) -> Diagram:
+    """Clean forward data-flow: the real input tensor flows top-to-bottom through each stage (in
+    call order, with the shape it produces) to the output tensor. Shapes come from one dummy
+    forward pass (meta, or a tiny CPU model). Used for codecs / conv backbones / detectors where
+    the data path — not an attention block — is the story (DAC, EnCodec, BLT, ResNet, ...)."""
+    from .introspect import _classify_mod
+
+    boxes: list[Box] = []
+    arrows: list[Arrow] = []
+    GX, GW = 300, 580  # flow column (facts panel sits at width-288)
+    y = PAD + 80
+    PFX = (am.config_class or "").replace("Config", "")
+    tree_by = {n["name"]: n for n in am.generic_tree}
+
+    def place(label, cls, subs, shape="op"):
+        nonlocal y
+        b = Box(GX, y, GW, 34 + 15 * len(subs), label, cls, sublabels=subs, shape=shape, small=True)
+        boxes.append(b)
+        y = b.y + b.h
+        return b
+
+    # stage list + shapes: prefer the captured forward (true call order + shapes); otherwise fall
+    # back to the structural top-level stages (definition order, no shapes) — always memory-free.
+    shapes = {s["name"]: s.get("out_shape") for s in am.flow}
+    if am.flow:
+        stages = [(s["name"], s["cls"]) for s in am.flow]
+        has_shapes = True
+    else:
+        stages = [(n["name"], n["cls"]) for n in am.generic_tree]
+        has_shapes = False
+
+    # input box
+    if am.flow_input:
+        in_label, in_sub = am.flow_input, []
+    else:
+        mods = am.modal_inputs or []
+        if "pixel_values" in mods or am.family == "image_classification":
+            in_label, in_sub = "input image", ["pixel_values [1, 3, H, W]"]
+        elif "input_features" in mods or am.family == "audio_classification":
+            in_label, in_sub = "input audio", ["input_values / input_features"]
+        else:
+            in_label, in_sub = "inputs", ["input_ids / pixel_values / input_values"]
+    prev = place(in_label, "c-io", in_sub, shape="io")
+
+    for name, cls in stages:
+        y += 24  # arrow gap
+        kind = _classify_mod(name, cls)
+        node = tree_by.get(name)
+        shp = shapes.get(name)
+        shape_lbl = f"   →  [{_shape_str(shp)}]" if shp else ""
+        hdr = f"{name}  ·  {_short_cls(cls, PFX)}{shape_lbl}"
+        sec_top = y
+        if node and node.get("children"):  # expose the stage's internal structure (×N blocks)
+            y += 28  # header strip
+            y = _emit_tree(boxes, node["children"], GX + 10, GW - 20, y, gp=7, prefix=PFX)
+            anchor = Box(GX, sec_top, GW, y - sec_top + 6, hdr, _NODE_CLS.get(kind, "c-layer"), shape="section")
+            boxes.append(anchor)
+            y += 6
+        else:
+            anchor = place(hdr, _NODE_CLS.get(kind, "c-layer"), [])
+        arrows.append(Arrow([(prev.x + prev.w / 2, prev.y + prev.h), (anchor.x + anchor.w / 2, sec_top)], cls="flow"))
+        prev = anchor
+    if am.flow_output:
+        y += 24
+        b = place(f"output  [{_shape_str(am.flow_output)}]", "c-io", [], shape="io")
+        arrows.append(Arrow([(prev.x + prev.w / 2, prev.y + prev.h), (b.x + b.w / 2, b.y)], cls="flow"))
+
+    facts = _facts(am)
+    height = max(y + PAD, PAD + 80 + len(facts) * 22 + 120)
+    sub_note = "forward data flow  (shapes from a dummy input pass)" if has_shapes else "staged data flow  (structure)"
+    return Diagram(
+        width=WF,
+        height=height,
+        boxes=boxes,
+        arrows=arrows,
+        title=am.model,
+        subtitle=f"{am.top_class or ''}  ·  {sub_note}",
+        legend=_STD_LEGEND,
+        facts=facts,
+        mode="full",
+        spine=False,
+    )
+
+
+# --------------------------------------------------------------- detailed full (Raschka)
+
+# geometry for the detailed view (wider canvas: schedule strip · column · masks · facts)
+WF = 1200
+STRIP_X = 40  # layer-schedule strip
+DCOL_X = 320  # left edge of central op column
+DCOL_W = 330  # width of central op column
+RAIL_X = 286  # residual rail x (left of the column)
+MASK_X = 700  # attention-mask grids column
+SEQ = 12  # dummy batch/seq for displayed tensor shapes: [1, 12]
+
+# layer_type -> color class (see render._STYLE)
+_LT_CLASS = {
+    "full_attention": "c-lt-full",
+    "sliding_attention": "c-lt-sliding",
+    "chunked_attention": "c-lt-chunked",
+    "compressed_sparse_attention": "c-lt-compressed",
+    "heavily_compressed_attention": "c-lt-heavy",
+    "linear_attention": "c-lt-linear",
+}
+
+
+def _lt_class(t: str) -> str:
+    if t in _LT_CLASS:
+        return _LT_CLASS[t]
+    n = t.lower()
+    if "linear" in n or "delta" in n or "gated" in n:
+        return "c-lt-linear"
+    if "mamba" in n or "ssm" in n:
+        return "c-lt-mamba"
+    if "sliding" in n:
+        return "c-lt-sliding"
+    if "chunk" in n:
+        return "c-lt-chunked"
+    if "compress" in n:
+        return "c-lt-compressed"
+    return "c-lt-full"
+
+
+def _dim(p) -> str:
+    if p is None:
+        return ""
+    return f"[{p.in_f or '?'}→{p.out_f or '?'}]"
+
+
+# module-tree node kind -> css class (for the recursive decomposition)
+# standardized node kind -> css class, shared by every view (attention internals, SSM mixer,
+# MoE, and the generic conv/codec/FFT module tree) so a given component is always the same colour.
+_NODE_CLS = {
+    "embedding": "c-embed",
+    "linear": "c-proj",
+    "conv": "c-conv",
+    "norm": "c-norm",
+    "act": "c-act",
+    "pool": "c-pool",
+    "dropout": "c-sub",
+    "rope": "c-rope",
+    "attention": "c-attn",
+    "mamba": "c-mamba",
+    "recurrent": "c-recur",
+    "quantizer": "c-quant",
+    "compressor": "c-lt-compressed",
+    "indexer": "c-lt-heavy",
+    "router": "c-moe",
+    "experts": "c-moe",
+    "head": "c-head",
+    "other": "c-sub",
+}
+
+# one standardized legend shared by all full-view layouts (only the colours that can appear)
+_STD_LEGEND = [
+    ("c-embed", "embedding"),
+    ("c-attn", "attention"),
+    ("c-mamba", "Mamba / SSM"),
+    ("c-moe", "MoE / experts"),
+    ("c-proj", "linear / proj"),
+    ("c-conv", "convolution"),
+    ("c-norm", "normalization"),
+    ("c-act", "activation"),
+    ("c-pool", "pooling"),
+    ("c-head", "head"),
+]
+
+
+def _short_cls(cls: str, prefix: str = "") -> str:
+    """Drop the model prefix and the redundant 'Activation' suffix: DeepseekV4RMSNorm→RMSNorm,
+    SiLUActivation→SiLU."""
+    s = cls
+    if prefix and s.startswith(prefix):
+        s = s[len(prefix) :]
+    if s.endswith("Activation") and len(s) > 10:
+        s = s[: -len("Activation")]
+    return s or cls
+
+
+def _leaf_label(nd, prefix: str) -> str:
+    """One compact line per leaf: name + dims; class shown only when it isn't a plain Linear."""
+    sc = _short_cls(nd["cls"], prefix)
+    dim = nd.get("dim")
+    if sc == "Linear":  # 'Linear' is implied by the [in→out] arrow
+        return f"{nd['name']}  {dim}" if dim else nd["name"]
+    return f"{nd['name']}  {sc}" + (f" {dim}" if dim else "")
+
+
+def _emit_tree(boxes, nodes, x, w, y, gp=7, depth=0, prefix=""):
+    """Recursively lay out a module tree: leaves are single-line chips, branches become titled
+    sub-sections containing their own children. Returns the y below the laid-out tree."""
+    i = 0
+    while i < len(nodes):
+        nd = nodes[i]
+        if nd.get("children"):  # a branch -> nested titled sub-section
+            sub_top = y
+            y += 25  # clear the white header strip
+            y = _emit_tree(boxes, nd["children"], x + 8, w - 16, y, gp, depth + 1, prefix)
+            label = f"{nd['name']} · {nd['cls']}" + (f"  ({nd['dim']})" if nd.get("dim") else "")
+            boxes.append(
+                Box(x, sub_top, w, y - sub_top + 6, label, _NODE_CLS.get(nd["kind"], "c-sub"), shape="section")
+            )
+            y += 6 + gp
+            i += 1
+        else:  # a run of leaf chips — content-sized: columns chosen by width, 2-line when needed
+            run = []
+            while i < len(nodes) and not nodes[i].get("children"):
+                run.append(nodes[i])
+                i += 1
+            # how many columns fit? each chip wants ~its text width; clamp 1..3 by available w
+            widest = max((text_width(_leaf_label(nd2, prefix), 12.5) for nd2 in run), default=80) + 18
+            per = max(1, min(3, int(w // max(widest, 120))))
+            for r in range(0, len(run), per):
+                rown = run[r : r + per]
+                cw = (w - (len(rown) - 1) * gp) / len(rown)
+                # a chip is 2 lines (name / class-dim) when the one-line label doesn't fit its cell
+                two = [text_width(_leaf_label(nd2, prefix), 12.5) > cw - 12 for nd2 in rown]
+                rh = 34 if any(two) else 22
+                for j, nd2 in enumerate(rown):
+                    ttl = f"{nd2['name']}: {nd2['cls']}" + (f" {nd2['dim']}" if nd2.get("dim") else "")
+                    if two[j]:  # split onto two lines instead of ellipsizing the dims away
+                        sc = _short_cls(nd2["cls"], prefix)
+                        sub = (
+                            (sc + (f" {nd2['dim']}" if nd2.get("dim") else ""))
+                            if sc != "Linear"
+                            else (nd2.get("dim") or "")
+                        )
+                        boxes.append(
+                            Box(
+                                int(x + j * (cw + gp)),
+                                y,
+                                int(cw),
+                                rh,
+                                nd2["name"],
+                                _NODE_CLS.get(nd2["kind"], "c-sub"),
+                                sublabels=[sub] if sub else [],
+                                shape="op",
+                                small=True,
+                                title=ttl,
+                            )
+                        )
+                    else:
+                        boxes.append(
+                            Box(
+                                int(x + j * (cw + gp)),
+                                y,
+                                int(cw),
+                                rh,
+                                _leaf_label(nd2, prefix),
+                                _NODE_CLS.get(nd2["kind"], "c-sub"),
+                                shape="op",
+                                small=True,
+                                title=ttl,
+                            )
+                        )
+                y += rh + gp
+    return y
+
+
+def _node_sig(node):
+    """Structural signature of a module node (class + dims + recursive children)."""
+    return (node.get("cls"), node.get("dim"), tuple(_node_sig(c) for c in (node.get("children") or [])))
+
+
+def _emit_merged(boxes, variants, types, x, w, y, gp=7, prefix=""):
+    """Render N aligned module trees, sharing identical sub-trees and branching side-by-side
+    ONLY where they actually differ. For DeepSeek-V4 attention this shows q/k/v/o once and the
+    differing `compressor` (HCA vs CSA+Indexer) side by side -- the real subtlety."""
+    base = variants[0]
+    names = [n["name"] for n in base]
+    if any([n["name"] for n in v] != names for v in variants):  # misaligned -> full side-by-side
+        nn = len(variants)
+        cw = (w - (nn - 1) * 18) / nn
+        mb = y
+        for k, v in enumerate(variants):
+            mb = max(mb, _emit_tree(boxes, v, int(x + k * (cw + 18)), int(cw), y, gp, prefix=prefix))
+        return mb
+    for idx in range(len(names)):
+        nodes = [v[idx] for v in variants]
+        if all(_node_sig(nd) == _node_sig(nodes[0]) for nd in nodes):  # identical -> render once
+            y = _emit_tree(boxes, [nodes[0]], x, w, y, gp, prefix=prefix)
+        else:  # diverges -> side by side, tagged by layer type
+            nn = len(nodes)
+            cw = (w - (nn - 1) * gp) / nn
+            top = y
+            mb = y
+            for k, nd in enumerate(nodes):
+                cx = int(x + k * (cw + gp))
+                boxes.append(Box(cx, top, int(cw), 18, f"▼ {types[k]}", "c-sub", shape="op", small=True))
+                mb = max(mb, _emit_tree(boxes, [nd], cx, int(cw), top + 22, gp, prefix=prefix))
+            y = mb
+    return y
+
+
+def _short_proj(name: str) -> str:
+    return name.replace("_proj", "").replace("_with_mqa", "·mqa").replace("kv_a", "kv↓").replace("kv_b", "kv↑")
+
+
+def _chips(boxes, items, x, w, y, per_row, cls="c-sub", gp=8, h=34):
+    """Place `items` = list of (label, sublabel) as a grid of small boxes; return new y."""
+    rows = [items[i : i + per_row] for i in range(0, len(items), per_row)]
+    for r in rows:
+        cw = (w - (len(r) - 1) * gp) / len(r)
+        for j, (lb, sub) in enumerate(r):
+            boxes.append(
+                Box(
+                    int(x + j * (cw + gp)),
+                    y,
+                    int(cw),
+                    h,
+                    lb,
+                    cls,
+                    sublabels=([sub] if sub else []),
+                    shape="op",
+                    small=True,
+                )
+            )
+        y += h + gp
+    return y
+
+
+def _shape3(h) -> str:
+    return f"[1, {SEQ}, {h or '?'}]"
+
+
+def build_full(am: ArchModel) -> Diagram:
+    """Dispatch: multimodal models get the multi-tower view, everything else the LLM column."""
+    if am.is_multimodal and am.towers and am.block is not None:
+        return build_multimodal(am)
+    # multi-stage pipelines (BLT, codecs, ...) read best as a data flow even when they use attention
+    if am.is_pipeline and am.generic_tree:
+        return _build_flow(am)
+    return _build_column(am)
+
+
+def _build_column(am: ArchModel) -> Diagram:
+    """Detailed, Raschka-style view, read **top-to-bottom like the code**: input_ids at the
+    top flow down through the embedding, the expanded transformer block (pre-norm →
+    self-attention with RoPE/Q/K/V/O and real dims → ⊕ residual → pre-norm → MLP or Sparse
+    MoE → ⊕ residual), final norm, LM head and softmax to the logits at the bottom. A
+    left-side strip shows the per-layer schedule (config.layer_types); attention-mask grids
+    for each distinct layer type are shown on the right. Tensor shapes use a dummy [1, 12]
+    input. Falls back to the compact view when block internals can't be extracted."""
+    if am.block is None or (am.block.attention is None and am.block.mixer is None and am.block.mlp is None):
+        # a real module tree → staged data-flow view (with shapes where a meta forward succeeded);
+        # only models that couldn't be built on meta fall back to the config-only schematic.
+        return _build_flow(am) if am.generic_tree else _build_compact(am)
+
+    blk = am.block
+    blk_is_ssm = blk.attention is None and blk.mixer is not None  # Mamba/SSM token-mixer block
+    GAP = 14
+    boxes: list[Box] = []
+    arrows: list[Arrow] = []
+    y = PAD + 80
+    H = am.hidden_size
+    PFX = (am.config_class or "").replace("Config", "")  # model prefix to strip from class names
+
+    def add(label, cls, subs=None, h=None, shape="op", glyph=None, w=DCOL_W, x=DCOL_X, title=None, grid=None):
+        nonlocal y
+        b = Box(
+            x,
+            y,
+            w,
+            h if h is not None else (30 + 15 * len(subs or [])),
+            label,
+            cls,
+            sublabels=subs or [],
+            shape=shape,
+            small=True,
+            glyph=glyph,
+            title=title,
+            grid=grid,
+        )
+        boxes.append(b)
+        y = y + b.h + GAP
+        return b
+
+    # which kind of model -> what the input / embedding / head look like
+    modality = (
+        "vision"
+        if am.family == "image_classification"
+        else ("audio" if am.family == "audio_classification" else "text")
+    )
+    enc = am.view == "encoder"
+    seq = len(am.tokens or [1] * 6) or 6
+
+    def _shapeN(h):
+        return f"[1, {seq}, {h or '?'}]"
+
+    # ---------- top: input ----------
+    if modality == "vision":
+        add("input image", "c-io", ["pixel_values [1, 3, H, W]"], h=32, shape="io")
+        add("Patch Embedding", "c-embed", [f"conv patches → {_shapeN(H)}", "+ position embeddings"], h=60)
+    elif modality == "audio":
+        add("input audio", "c-io", ["input_features [1, n_mels, T]"], h=32, shape="io")
+        add("Feature Projection", "c-embed", [f"conv/linear → {_shapeN(H)}", "+ position embeddings"], h=60)
+    else:
+        toks = am.tokens or ["Hey", ",", "␣how", "␣are", "␣you", "?"]
+        seq = len(toks)
+        add("input_ids", "c-io", [f'tokenize("Hey, how are you?")  →  [1, {seq}]'], h=32, shape="io")
+        tgp = 5
+        tw = (DCOL_W - (seq - 1) * tgp) / seq
+        for j, t in enumerate(toks):
+            label = t if len(t) <= 6 else t[:6] + "…"
+            boxes.append(
+                Box(
+                    int(DCOL_X + j * (tw + tgp)),
+                    y,
+                    max(14, int(tw)),
+                    26,
+                    label,
+                    "c-sub",
+                    sublabels=[str(j)],
+                    shape="op",
+                    small=True,
+                    title=f"token {j}: {t!r}",
+                )
+            )
+        y += 26 + 18
+        pos_note = [f"+ {am.positional} positions"] if (am.positional and am.positional not in ("RoPE", "n/a")) else []
+        add(
+            "Token Embedding",
+            "c-embed",
+            [f"weight [{am.vocab_size or '?'} × {H or '?'}]   →   {_shapeN(H)}"] + pos_note,
+            h=30 + 15 * (1 + len(pos_note)),
+        )
+    body_top = boxes[-1].y  # the embedding box top — start of the "<Model>" body wrapper
+
+    # inner-chip geometry shared by the attention / MoE / MLP sections
+    IX, IW, GP = DCOL_X + 12, DCOL_W - 24, 8
+
+    def section(label, cls, body, badge=None):
+        """Draw a filled section panel containing inner chips produced by ``body(y)``."""
+        nonlocal y
+        sec_top = y
+        y += 26  # header room
+        y = body(y)
+        sec = Box(DCOL_X, sec_top, DCOL_W, y - sec_top + 6, label, cls, shape="section", badge=badge)
+        boxes.append(sec)
+        y = y + 6 + GAP
+        return sec
+
+    # ---------- transformer block ----------
+    y += 22  # headroom for the dashed container's "decoder block ×N" label
+    block_top = y
+    entry_y = block_top  # residual stream entry
+    add(
+        f"{blk.pre_attn_norm or 'Norm'}",
+        "c-norm",
+        [("pre-mixer  " if blk_is_ssm else "pre-attention  ") + _shapeN(H)],
+        h=28,
+    )
+
+    # ----- self-attention: ONE section; shared sub-modules rendered once, and only the parts
+    # ----- that actually differ across layer types (e.g. the compressor) branch side by side.
+    a = blk.attention
+    from .introspect import _short_layer_type
+
+    attn_right = DCOL_X + DCOL_W
+    variants = am.attention_variants
+    if variants:
+        n = min(len(variants), 3)
+        types = [_short_layer_type(v["type"]) for v in variants[:n]]
+        sc_by_type = {sc["type"]: sc for sc in am.sparse_components}
+        diverges = n > 1
+        Wsec = DCOL_W if not diverges else max(DCOL_W, n * 300)
+        gqa = ""
+        if a and a.n_heads and a.n_kv and a.n_kv < a.n_heads:
+            gqa = f" · GQA {a.n_heads}:{a.n_kv}"
+        elif a and a.n_kv == 1:
+            gqa = " · MQA"
+        vtop = y
+        vy = vtop + 26
+        if n == 1:
+            vy = _emit_tree(boxes, variants[0]["children"], DCOL_X + 10, Wsec - 20, vy, gp=GP)
+        else:
+            vy = _emit_merged(boxes, [v["children"] for v in variants[:n]], types, DCOL_X + 10, Wsec - 20, vy, gp=GP)
+        # core op footer (shared): SDPA for attention, selective scan for SSM/Mamba mixers
+        is_mixer = a is None and blk.mixer is not None
+        boxes.append(
+            Box(
+                DCOL_X + 10,
+                vy,
+                Wsec - 20,
+                30,
+                "selective state-space scan (SSM)" if is_mixer else "scaled dot-product attention",
+                "c-sub",
+                sublabels=([f"{a.n_heads or '?'} heads · head_dim {a.head_dim or '?'}{gqa}"] if a else []),
+                shape="op",
+                small=True,
+            )
+        )
+        vy += 30 + GP
+        # masks: one per layer type, side by side (this is where the compressors differ in effect)
+        mask_items = [
+            (v["type"], (sc_by_type.get(v["type"], {}).get("mask") or am.attn_patterns.get(v["type"])))
+            for v in variants[:n]
+        ]
+        mask_items = [(t, g) for t, g in mask_items if g]
+        if not mask_items:  # encoders attend bidirectionally; decoders are causal
+            sq = 14
+            if enc:
+                grid = [[1] * sq for _ in range(sq)]
+                mask_items = [("bidirectional", grid)]
+            else:
+                grid = [[1 if j <= i else 0 for j in range(sq)] for i in range(sq)]
+                mask_items = [("causal", grid)]
+        if mask_items:
+            mw = (Wsec - 20 - (len(mask_items) - 1) * 16) / len(mask_items)
+            mtop = vy + 4
+            mb = mtop
+            for k, (lt, grid) in enumerate(mask_items):
+                sc = sc_by_type.get(lt)
+                cols = len(grid[0])
+                cellp = max(3, min(10, int((mw - 6) // cols)))
+                gx = int(DCOL_X + 10 + k * (mw + 16) + (mw - cellp * cols) / 2)
+                cap = (
+                    f"{_short_layer_type(lt)} mask  m={sc['display_m']}·{sc['n_comp']}c"
+                    if sc
+                    else f"{_short_layer_type(lt)} mask"
+                )
+                boxes.append(
+                    Box(
+                        gx,
+                        mtop + 6,
+                        cellp * cols,
+                        cellp * len(grid),
+                        cap,
+                        _lt_class(lt),
+                        shape="grid",
+                        grid=grid,
+                        grid_split=(sc.get("mask_split") if sc else None),
+                        title=f"{lt} attention mask (q↓ × k→)",
+                    )
+                )
+                mb = max(mb, mtop + 6 + cellp * len(grid))
+            vy = mb
+        note = ""
+        if diverges:
+            sigs = [tuple(_node_sig(c) for c in v["children"]) for v in variants[:n]]
+            struct_diff = any(s != sigs[0] for s in sigs)
+            note = (
+                f"  ({n} layer types — modules differ ▼)"
+                if struct_diff
+                else f"  ({n} layer types — same modules, mask differs)"
+            )
+        if is_mixer:
+            hdr = f"Token Mixer (SSM) · {blk.mixer.cls}{note}"
+        else:
+            hdr = f"Self-Attention · {a.cls if a else 'attention'}{note}"
+        boxes.append(Box(DCOL_X, vtop, Wsec, vy - vtop + 6, hdr, "c-attn", shape="section"))
+        attn_right = DCOL_X + Wsec
+        y = vy + 6 + GAP
+        if a and a.rope:
+            ry = vtop + 30
+            rope_label = f"RoPE  θ={int(am.rope_theta):,}" if am.rope_theta else "RoPE"
+            rnode = Box(
+                DCOL_X - 168,
+                ry - 4,
+                132,
+                40,
+                rope_label,
+                "c-rope",
+                sublabels=["rotary positions → Q, K"],
+                shape="op",
+                small=True,
+            )
+            boxes.append(rnode)
+            arrows.append(Arrow([(rnode.x + rnode.w, ry + 14), (DCOL_X, ry + 14)], cls="rope"))
+    else:
+        add("Token Mixing  (linear-attention / SSM)", "c-linattn", [blk.layer_class], h=44)
+    plus_attn = add("", "c-add", h=26, shape="circle", glyph="+", title="residual add")
+
+    # ----- MLP / Sparse MoE: FULL recursive decomposition, variants side by side -----
+    # (Mamba/SSM blocks are norm → mixer → residual with NO feed-forward, so skip the FFN entirely)
+    mm = blk.mlp
+    mvars = am.mlp_variants if len(am.mlp_variants) > 1 else None
+    has_mlp = bool(mm or (am.mlp_tree and am.mlp_tree.get("children")) or mvars)
+    if has_mlp:
+        add(f"{blk.post_attn_norm or 'Norm'}", "c-norm", ["pre-FFN  " + _shapeN(H)], h=28)
+    if mvars:
+        # distinct FFN types (e.g. DeepSeek-V4 moe / hash_moe): shared sub-modules once, only
+        # the differing part (the router) branches side by side.
+        nm = min(len(mvars), 3)
+        Wm = max(DCOL_W, nm * 300)
+        mtypes = [mv["type"] for mv in mvars[:nm]]
+        vtop = y
+        my = vtop + 26
+        my = _emit_merged(boxes, [mv["children"] for mv in mvars[:nm]], mtypes, DCOL_X + 10, Wm - 20, my, gp=GP)
+        cls0 = mvars[0]["cls"]
+        boxes.append(
+            Box(
+                DCOL_X,
+                vtop,
+                Wm,
+                my - vtop + 6,
+                f"Sparse MoE · {cls0}  ({nm} FFN types ▼)",
+                "c-moe",
+                shape="section",
+                badge="MoE",
+            )
+        )
+        attn_right = max(attn_right, DCOL_X + Wm)
+        y = my + 6 + GAP
+    elif am.mlp_tree and am.mlp_tree.get("children"):
+        cls = am.mlp_tree["cls"]
+        is_moe = bool(mm and mm.is_moe)
+        section(
+            ("Sparse MoE · " if is_moe else "MLP · ") + cls,
+            "c-moe" if is_moe else "c-mlp",
+            lambda yy: _emit_tree(boxes, am.mlp_tree["children"], IX, IW, yy, gp=GP),
+            badge="MoE" if is_moe else None,
+        )
+    elif mm:
+
+        def mlp_body(y):
+            items = []
+            if mm.gate:
+                items.append(("gate_proj", _dim(mm.gate)))
+            if mm.up:
+                items.append(("up_proj", _dim(mm.up)))
+            per = len(items) if items else 1
+            if items:
+                y = _chips(boxes, items, IX, IW, y, per, gp=GP, h=32)
+            boxes.append(
+                Box(IX, y, IW, 28, f"{mm.act}(gate) ⊙ up" if mm.gate else mm.act, "c-sub", shape="op", small=True)
+            )
+            y += 28 + GP
+            boxes.append(Box(IX, y, IW, 32, f"down_proj  {_dim(mm.down)}", "c-sub", shape="op", small=True))
+            y += 32
+            return y
+
+        section(f"MLP · {mm.cls}", "c-mlp", mlp_body)
+    elif has_mlp:
+        add("MLP", "c-mlp", h=30)
+    # second residual add only exists when there is an FFN; SSM blocks have a single residual
+    plus_mlp = add("", "c-add", h=26, shape="circle", glyph="+", title="residual add") if has_mlp else plus_attn
+    block_bot = y - GAP
+
+    container_w = max(DCOL_W, attn_right - DCOL_X) + 52
+    container = Box(
+        DCOL_X - 26,
+        block_top - 22,
+        container_w,
+        block_bot - block_top + 30,
+        blk.layer_class,  # the real layer class, e.g. LlamaDecoderLayer / ASTLayer
+        "c-block",
+        shape="container",
+        badge=f"× {am.num_layers or '?'}",
+        title=f"{am.num_layers} × {blk.layer_class}",
+    )
+
+    # ---------- bottom: final norm, then the task head (depends on family) ----------
+    add(f"Final {am.norm_type or 'Norm'}", "c-norm", [_shapeN(H)], h=28)
+    body_bot = boxes[-1].y + boxes[-1].h  # end of the "<Model>" body (everything above = base model)
+    head_title = am.head_class  # e.g. ViTForImageClassification / GemmaForCausalLM (the For wrapper)
+    if enc:  # encoder / classifier head
+        nl = am.num_labels or "?"
+        add("Pool  (CLS / mean)", "c-head", [f"{_shapeN(H)} → [1, {H or '?'}]"], h=40, title=head_title)
+        add("Classifier head", "c-head", [f"Linear [{H or '?'}→{nl}]"], h=40, title=head_title)
+        add("Softmax", "c-soft", h=26)
+        add("class logits", "c-io", [f"[1, {nl}]  ({nl} classes)"], h=30, shape="io")
+    else:  # decoder LM head
+        head_sub = [f"Linear [{H or '?'}→{am.vocab_size or '?'}]" + ("  · tied" if am.tie_word_embeddings else "")]
+        add("LM Head", "c-head", head_sub, h=42, title=head_title)
+        add("Softmax", "c-soft", h=26)
+        add("logits", "c-io", [f"[1, {seq}, {am.vocab_size or '?'}]"], h=30, shape="io")
+
+    # ---------- wrapper hierarchy: base <Model> body and <ForXxx> task head, STACKED (adjacent,
+    # not nested) so the two labels never overlap ----------
+    ww = (attn_right - DCOL_X) + 68
+    if am.top_class:
+        boxes.append(
+            Box(
+                DCOL_X - 34,
+                body_top - 14,
+                ww,
+                body_bot - body_top + 20,
+                f"{am.top_class}  ·  base model",
+                "c-block",
+                shape="container",
+            )
+        )
+    if head_title:
+        boxes.append(
+            Box(
+                DCOL_X - 34,
+                body_bot + 6,
+                ww,
+                (y - GAP) - (body_bot + 6) + 8,
+                f"{head_title}  ·  task head",
+                "c-block",
+                shape="container",
+            )
+        )
+
+    # ---------- residual skip arrows (flow downward) ----------
+    def midy(b):
+        return b.y + b.h // 2
+
+    arrows.append(
+        Arrow(
+            [(DCOL_X, entry_y), (RAIL_X, entry_y), (RAIL_X, midy(plus_attn)), (plus_attn.x, midy(plus_attn))],
+            cls="residual",
+            dashed=True,
+        )
+    )
+    if has_mlp:
+        arrows.append(
+            Arrow(
+                [
+                    (plus_attn.x, midy(plus_attn)),
+                    (RAIL_X - 16, midy(plus_attn)),
+                    (RAIL_X - 16, midy(plus_mlp)),
+                    (plus_mlp.x, midy(plus_mlp)),
+                ],
+                cls="residual",
+                dashed=True,
+            )
+        )
+
+    boxes.insert(0, container)
+
+    # ---------- layer-schedule strip (config.layer_types) — one labelled cell per layer ----
+    legend = [("c-attn", "self-attention")]
+    strip_bottom = block_bot
+    # only show the per-layer schedule when layers actually differ (≥2 distinct types)
+    if am.layer_types and len(set(am.layer_types)) >= 2:
+        from .introspect import _short_layer_type
+
+        n = len(am.layer_types)
+        SW = 46  # strip width
+        cell_h = max(15, (block_bot - block_top) / n)  # at least 15px so the index fits
+        boxes.append(Box(STRIP_X - 4, block_top - 24, SW + 12, 16, "layers (idx)", "c-block", shape="container"))
+        for i, lt in enumerate(am.layer_types):
+            cy = int(block_top + i * cell_h)
+            boxes.append(
+                Box(
+                    STRIP_X,
+                    cy,
+                    SW,
+                    max(13, int(cell_h) - 2),
+                    str(i),
+                    _lt_class(lt),
+                    shape="cell",
+                    title=f"layer {i}: {lt}",
+                )
+            )
+        strip_bottom = int(block_top + n * cell_h)
+        seen = []
+        for lt in am.layer_types:
+            if lt not in seen:
+                seen.append(lt)
+        for lt in seen:
+            legend.append((_lt_class(lt), _short_layer_type(lt)))
+
+    # (attention-mask grids are now drawn INSIDE each variant card, so no separate column)
+    gy = 0
+    legend += [
+        ("c-proj", "linear / proj"),
+        ("c-moe", "MoE / experts") if am.is_moe else ("c-mlp", "MLP / feed-forward"),
+        ("c-rope", "RoPE"),
+        ("c-norm", "normalization"),
+        ("residual", "residual / skip"),
+    ]
+
+    facts = _facts(am)
+    mask_bottom = gy
+    height = max(y + PAD, strip_bottom + PAD, mask_bottom + PAD, PAD + 80 + len(facts) * 22 + len(legend) * 18 + 120)
+    sub = f"{am.decoder_kind} · {am.layer_summary or ''}"
+    width = max(WF, attn_right + 320)  # widen for side-by-side attention variants + facts panel
+    return Diagram(
+        width=width,
+        height=height,
+        boxes=boxes,
+        arrows=arrows,
+        title=am.model,
+        subtitle=sub,
+        legend=legend,
+        facts=facts,
+        mode="full",
+        spine=True,
+    )
+
+
+# --------------------------------------------------------- multimodal multi-tower full view
+
+
+def build_multimodal(am: ArchModel) -> Diagram:
+    """VLM / audio view: encoder tower(s) → projector → the LLM (the detailed column).
+
+    The LLM column is built by ``_build_column`` and shifted right; the modality towers are
+    drawn in a left lane, with a fusion arrow into the LLM (merge at modality tokens), or
+    cross-attention edges into specific LLM layers when the text config declares them.
+    """
+    d = _build_column(am)
+    dx = 300
+    for b in d.boxes:
+        b.x += dx
+    for a in d.arrows:
+        a.points = [(x + dx, y) for (x, y) in a.points]
+    d.width += dx
+
+    emb = next((b for b in d.boxes if b.cls == "c-embed"), None)
+    container = next((b for b in d.boxes if b.shape == "container" and b.cls == "c-block"), None)
+
+    TX, TW = 36, 250
+    ty = PAD + 92
+    encoders = [t for t in am.towers if t["role"] in ("vision", "audio")]
+    proj = next((t for t in am.towers if t["role"] == "projector"), None)
+
+    mmpfx = (am.config_class or "").replace("Config", "")
+
+    def tb(label, cls, subs, h, badge=None, shape="op"):
+        nonlocal ty
+        b = Box(TX, ty, TW, h, label, cls, sublabels=subs or [], shape=shape, small=True, badge=badge)
+        d.boxes.append(b)
+        ty += h + 16
+        return b
+
+    def tower(t, label, cls, summary, badge=None):
+        """Render a tower: fully decomposed when small; a representative block ×N for big
+        encoder stacks; else a summary box. So every tower shows what's inside it."""
+        nonlocal ty
+        ch = t.get("children")
+        if not ch and t.get("block_children"):
+            # big encoder: outer section + one representative layer (decomposed) with ×N
+            top = ty
+            inner_top = top + 26
+            blabel = f"{t['block_class']}"
+            by = _emit_tree(d.boxes, t["block_children"], TX + 18, TW - 36, inner_top + 22, gp=6, prefix=mmpfx)
+            d.boxes.append(
+                Box(
+                    TX + 10,
+                    inner_top,
+                    TW - 20,
+                    by - inner_top + 6,
+                    blabel,
+                    cls,
+                    shape="section",
+                    badge=f"×{t['block_n']}",
+                )
+            )
+            b = Box(TX, top, TW, (by + 6) - top + 6, label, cls, shape="section", badge=badge)
+            d.boxes.append(b)
+            ty = by + 6 + 6 + 16
+            return b
+        if ch:
+            top = ty
+            yy = _emit_tree(d.boxes, ch, TX + 10, TW - 20, top + 26, gp=6, prefix=mmpfx)
+            b = Box(TX, top, TW, yy - top + 6, label, cls, shape="section", badge=badge)
+            d.boxes.append(b)
+            ty = yy + 6 + 16
+            return b
+        return tb(label, cls, summary, 52, badge=badge)
+
+    enc_boxes = []
+    for t in encoders:
+        if t["role"] == "vision":
+            inp, ishape, cls, role = "pixel_values", "[1, 3, 336, 336]", "c-vision", "Vision encoder"
+            img = Box(TX + (TW - 76) // 2, ty, 76, 64, "example image", "c-vision", shape="image")
+            d.boxes.append(img)
+            ty += 64 + 22
+        else:
+            inp, ishape, cls, role = "input_features", "[1, 128, 3000]", "c-audio", "Audio encoder"
+            img = None
+        ib = tb(inp, "c-io", [ishape], 30, shape="io")
+        if img is not None:
+            d.arrows.append(Arrow([(img.x + img.w // 2, img.y + img.h), (ib.x + ib.w // 2, ib.y)], cls="flow"))
+        cls_line = t["cls"] + ("  · via AutoModel" if am.auto_classes else "")
+        eb = tower(
+            t,
+            f"{role} · {t['cls']}",
+            cls,
+            [cls_line, f"{t['model_type'] or ''} · {t['layers'] or '?'} layers · h={t['hidden'] or '?'}"],
+            badge=(f"×{t['layers']}" if t["layers"] else None),
+        )
+        d.arrows.append(Arrow([(ib.x + ib.w // 2, ib.y + ib.h), (eb.x + eb.w // 2, eb.y)], cls="flow"))
+        enc_boxes.append(eb)
+
+    pj = tower(
+        proj or {}, f"Projector · {proj['cls'] if proj else 'Linear'}", "c-proj", ["align features → text hidden dim"]
+    )
+    for eb in enc_boxes:
+        d.arrows.append(Arrow([(eb.x + eb.w // 2, eb.y + eb.h), (pj.x + pj.w // 2, pj.y)], cls="flow"))
+
+    # fusion into the LLM
+    px = pj.x + pj.w
+    pmid = pj.y + pj.h // 2
+    if am.cross_attention_layers and container is not None:
+        n = len(am.cross_attention_layers)
+        target_x, target_y = container.x, container.y + container.h // 2
+        d.arrows.append(
+            Arrow(
+                [(px, pmid), ((px + target_x) // 2, pmid), ((px + target_x) // 2, target_y), (target_x, target_y)],
+                cls="xattn",
+            )
+        )
+        d.boxes.append(
+            Box(
+                int((px + target_x) // 2) - 70,
+                target_y - 34,
+                150,
+                18,
+                f"cross-attention @ {n} layers",
+                "c-proj",
+                shape="op",
+                small=True,
+            )
+        )
+    elif emb is not None:
+        target_x, target_y = emb.x, emb.y + emb.h // 2
+        d.arrows.append(
+            Arrow(
+                [(px, pmid), ((px + target_x) // 2, pmid), ((px + target_x) // 2, target_y), (target_x, target_y)],
+                cls="flow",
+            )
+        )
+        d.boxes.append(
+            Box(
+                int((px + target_x) // 2) - 78,
+                target_y - 30,
+                168,
+                18,
+                "merge at modality tokens",
+                "c-proj",
+                shape="op",
+                small=True,
+            )
+        )
+
+    # image ↔ text attention mask (the cross-modal mask the LLM actually uses)
+    if any(t["role"] == "vision" for t in encoders):
+        from .masks import image_text_mask
+
+        n_img, n_text = 5, 7
+        grid, split = image_text_mask(n_img, n_text, image_bidirectional=am.image_bidirectional)
+        cellp = 13
+        ty += 16  # gap below the projector
+        mx = TX + (TW - cellp * len(grid)) // 2
+        my = ty + 22
+        kind = "prefix-LM: image bidirectional" if am.image_bidirectional else "fully causal"
+        d.boxes.append(Box(TX - 4, ty, TW + 8, 16, "image ⊕ text attention", "c-block", shape="container"))
+        d.boxes.append(
+            Box(
+                mx,
+                my,
+                cellp * len(grid),
+                cellp * len(grid),
+                f"{n_img} img + {n_text} text · {kind}",
+                "c-vision",
+                shape="grid",
+                grid=grid,
+                grid_split=split,
+                title="image↓text query × image|text key",
+            )
+        )
+        ty = my + cellp * len(grid) + 16
+
+    # legend additions
+    if any(t["role"] == "vision" for t in encoders):
+        d.legend.insert(0, ("c-vision", "vision encoder"))
+    if any(t["role"] == "audio" for t in encoders):
+        d.legend.insert(0, ("c-audio", "audio encoder"))
+    d.legend.append(("c-proj", "projector / fusion"))
+    if am.cross_attention_layers:
+        d.legend.append(("xattn", "cross-attention"))
+
+    d.subtitle = f"multimodal · {'+'.join(t['role'] for t in encoders)} → projector → LLM · {am.decoder_kind}"
+    d.facts.insert(1, ("inputs", ", ".join(am.modal_inputs)))
+    if am.auto_classes:
+        d.facts.insert(2, ("built via", ", ".join(am.auto_classes[:2])))
+    d.height = max(d.height, ty + PAD)
+    return d
+
+
+# ----------------------------------------------------------------------------- diff mode
+
+
+# map a modular class name to which architectural slot it touches.
+# We match on SUFFIXES, not substrings: model names themselves often contain component
+# keywords (e.g. "Qwen2Moe...") which would otherwise mis-route every class to "moe".
+def _slot_for_class(name: str) -> str:
+    n = name.lower()
+    if n.endswith("config"):
+        return "config"
+    if n.endswith("rotaryembedding") or n.endswith("rotaryembeddings"):
+        return "rope"
+    if n.endswith(("router", "experts", "expert", "moeblock", "sparsemoeblock", "moe", "moemlp")):
+        return "moe"
+    if n.endswith(("mlp", "feedforward", "ffn")):
+        return "mlp"
+    if n.endswith(("attention", "attn", "sdpaattention", "flashattention2")):
+        return "attention"
+    if n.endswith(("rmsnorm", "layernorm", "norm")):
+        return "norm"
+    if (
+        n.endswith(("pretrainedmodel",))
+        or "forcausallm" in n
+        or "forsequence" in n
+        or "fortoken" in n
+        or "forquestion" in n
+        or "forconditional" in n
+        or "forretrieval" in n
+        or "withlmhead" in n
+    ):
+        return "head"
+    if n.endswith(("decoderlayer", "encoderlayer", "layer", "block", "mixer")):
+        return "layer"
+    if n.endswith(("embedding", "embeddings")):
+        return "embedding"
+    if n.endswith("model"):
+        return "layer"  # the decoder/model wrapper -> the stack
+    return "layer"
+
+
+def _change_for(cc: ClassChange) -> str | None:
+    if cc.relation == "new":
+        return "added"
+    if cc.deleted_methods or cc.deleted_attrs:
+        return "deleted"
+    if cc.n_changes == 0:
+        return None  # trivial / unchanged
+    if cc.added_methods or cc.added_attrs:
+        return "added"
+    return "overridden"
+
+
+def _change_detail(cc: ClassChange) -> str:
+    bits = []
+    if cc.overridden_methods:
+        bits.append("ovr " + ",".join(cc.overridden_methods))
+    if cc.added_methods:
+        bits.append("add " + ",".join(cc.added_methods))
+    if cc.deleted_methods:
+        bits.append("del " + ",".join(cc.deleted_methods))
+    na = len(cc.added_attrs) + len(cc.overridden_attrs)
+    if na:
+        bits.append(f"{na} attr")
+    return "; ".join(bits) if bits else "rename only"
+
+
+def build_diff(am: ArchModel, ad: ArchDiff) -> Diagram:
+    """Diff over the *real architecture*: render the full detailed view, then colour each
+    class-box by how it relates to the parent it inherits from -- added (green), overridden
+    (amber), deleted (red). Everything inherited unchanged / copy-pasted is greyed (ghosted),
+    so at a glance you see what the model is made of AND exactly what it changed."""
+    d = build_full(am)
+    d.mode = "diff"
+    d.title = am.model
+
+    if not ad.is_modular:
+        # standalone: nothing inherited -> show the full architecture, note it
+        d.subtitle = (ad.note or "standalone (no modular parent)") + " · shown in full"
+        return d
+
+    # change type + detail per modular class (skip trivial → treated as inherited/ghost).
+    # Box colour: a class redefined in the modular file (even a "new" subclass like
+    # GemmaRMSNorm, which reimplements the inherited norm) is a CHANGE -> amber, not green.
+    # Green is reserved for genuinely net-new *submodules* (added below from the parent diff).
+    by_cls: dict[str, str] = {}
+    detail: dict[str, ClassChange] = {}
+    changes_list: list[tuple[str, str, str]] = []
+    rank = {"deleted": 0, "overridden": 1, "added": 2}
+    for cc in ad.changes:
+        ch = _change_for(cc)
+        detail[cc.name] = cc
+        if ch is not None:
+            by_cls[cc.name] = "deleted" if ch == "deleted" else "overridden"
+            changes_list.append((ch, cc.name, _change_detail(cc)))
+    changes_list.sort(key=lambda t: (rank.get(t[0], 9), t[1]))
+
+    # longest class names first so we match e.g. GemmaAttention before Gemma
+    changed_names = sorted(by_cls, key=len, reverse=True)
+
+    def match(text: str) -> str | None:
+        for nm in changed_names:
+            if nm in text:
+                return by_cls[nm]
+        return None
+
+    def by_keyword(*kws: str) -> str | None:
+        # the embedding/head boxes have generic labels; match their changed class by keyword
+        for nm in changed_names:
+            low = nm.lower()
+            if any(k in low for k in kws):
+                return by_cls[nm]
+        return None
+
+    # colour each box; ghost everything that didn't change (inherited / copy-pasted)
+    for b in d.boxes:
+        if b.shape in ("container", "io", "circle", "cell", "grid"):
+            continue
+        ch = match(f"{b.label or ''}  {b.title or ''}")
+        if ch is None and b.cls == "c-embed":
+            ch = by_keyword("embedding", "embed")
+        if ch is None and b.cls == "c-head":
+            ch = by_keyword("for", "lmhead", "classifier", "classification", "head")
+        if ch:
+            b.change = ch
+            b.ghost = False
+        else:
+            b.ghost = True  # inherited unchanged / not a (changed) modular class
+
+    # compare against the PARENT architecture itself, so overridden blocks reveal WHAT changed
+    pam = _parent_arch(ad.parent_model) if ad.parent_model else None
+    if pam is not None:
+        pnames: set = set()
+        for v in pam.attention_variants:
+            _collect_submodule_names(v["children"], pnames)
+        if pam.mlp_tree:
+            _collect_submodule_names(pam.mlp_tree.get("children") or [], pnames)
+        # a submodule present in the child but NOT in the parent is an addition (e.g. qwen3's
+        # q_norm / k_norm vs qwen2) -> light it green even inside an "overridden" block
+        if pnames:
+            for b in d.boxes:
+                if b.shape == "op" and b.title and ": " in b.title and not b.change:
+                    nm, _, rest = b.title.partition(": ")
+                    # only real submodule chips (title "name: ClassName"), not token/io chips
+                    if not rest[:1].isupper():
+                        continue
+                    if nm and nm not in pnames:
+                        b.change = "added"
+                        b.ghost = False
+        # inherited (ghosted) section headers: show the PARENT class they came from
+        cpfx = (am.config_class or "").replace("Config", "")
+        ppfx = (pam.config_class or "").replace("Config", "")
+        if cpfx and ppfx and cpfx != ppfx:
+            for b in d.boxes:
+                if b.ghost and b.shape == "section" and cpfx in (b.label or ""):
+                    b.label = b.label.replace(cpfx, ppfx) + "  ↩ inherited"
+
+    t = ad.totals
+    extra = f" (+{','.join(p for p in ad.parent_models if p != ad.parent_model)})" if len(ad.parent_models) > 1 else ""
+    d.subtitle = (
+        f"diff vs {ad.parent_model}{extra}  ·  {t['overridden']} overridden · {t['added']} added · "
+        f"{t['deleted']} deleted · {t['new_classes']} new · {t['trivial']} inherited-as-is"
+    )
+    d.changes = changes_list
+    d.legend = [
+        ("ch-added", "new submodule (vs parent)"),
+        ("ch-over", "changed / redefined"),
+        ("ch-deleted", "deleted"),
+        ("ghost", "inherited / copy-pasted"),
+    ]
+    d.facts = [
+        ("model id", (am.checkpoint or am.model)[:30]),
+        ("parent", ad.parent_model or "—"),
+        ("classes", str(len(ad.changes))),
+        ("overridden", str(t["overridden"])),
+        ("added", str(t["added"])),
+        ("new classes", str(t["new_classes"])),
+        ("inherited as-is", str(t["trivial"])),
+    ]
+    # ensure height fits the changes panel
+    panel_h = len(d.facts) * 22 + len(d.legend) * 18 + len(d.changes) * 28 + 200
+    d.height = max(d.height, PAD + 80 + panel_h)
+    return d

arch_svg/masks.py

@@ -0,0 +1,149 @@
+"""Attention-mask *pattern* per layer type, generated offline.
+
+The goal is to make the per-layer attention pattern **visible** (full-causal vs sliding vs
+chunked vs compressed vs bidirectional). Real masks built from config are technically exact,
+but production sliding windows (128, 4096, ...) are far larger than any sequence we can draw,
+so they collapse to a plain causal triangle and become indistinguishable. We therefore draw
+**schematic** patterns on a small ``seq`` grid using an illustrative window/chunk, and label
+the *true* window size elsewhere. Where a real pattern is small enough to be faithful (window
+< seq) the schematic coincides with it. Never raises.
+"""
+
+from __future__ import annotations
+
+
+_WIN = 6  # illustrative sliding window / chunk size on the schematic grid
+
+
+def _causal(seq: int) -> list[list[int]]:
+    return [[1 if j <= i else 0 for j in range(seq)] for i in range(seq)]
+
+
+def _full(seq: int) -> list[list[int]]:
+    return [[1] * seq for _ in range(seq)]
+
+
+def _sliding(seq: int, win: int) -> list[list[int]]:
+    return [[1 if (0 <= i - j < win) else 0 for j in range(seq)] for i in range(seq)]
+
+
+def _chunked(seq: int, chunk: int) -> list[list[int]]:
+    # causal within the current chunk only
+    return [[1 if (j <= i and j // chunk == i // chunk) else 0 for j in range(seq)] for i in range(seq)]
+
+
+def _compressed(seq: int, win: int, stride: int = 2) -> list[list[int]]:
+    # local sliding window + sparse/strided long-range keys (schematic of compressed attention)
+    g = _sliding(seq, win)
+    for i in range(seq):
+        for j in range(0, i - win + 1, stride):
+            g[i][j] = 1
+    return g
+
+
+def _compressed_kv(seq: int, real_rate: int, topk: int | None = None) -> list[list[int]]:
+    """Schematic of attention over a *compressed KV cache*: queries (rows) attend to a
+    compressed key sequence (cols). Columns = seq / display-rate (the real rate is usually too
+    large to draw, so we cap it for visibility and label the true rate elsewhere). Causal over
+    the compressed positions; for the sparse variant only the most-recent ``topk`` compressed
+    keys are kept (a band)."""
+    disp = max(2, min(real_rate, 6))
+    cols = max(2, -(-seq // disp))  # ceil
+    grid = []
+    for i in range(seq):
+        ci = i // disp  # compressed position the query has reached
+        row = [1 if c <= ci else 0 for c in range(cols)]
+        if topk:  # sparse: keep only the most-recent compressed keys (schematic band/top-k)
+            for c in range(cols):
+                if row[c] and (ci - c) >= max(2, cols // 2):
+                    row[c] = 0
+        grid.append(row)
+    return grid
+
+
+# a fixed example sentence so mask figures read like the reference (token-labelled axes)
+EXAMPLE_WORDS = [
+    "The",
+    "quick",
+    "brown",
+    "fox",
+    "jumps",
+    "over",
+    "the",
+    "lazy",
+    "dog",
+    "ate",
+    "a",
+    "small",
+    "red",
+    "fish",
+    "this",
+    "morning",
+]
+
+
+def concat_compressed_mask(seq: int, m: int, topk: int | None = None):
+    """Reproduce the *actual* mask DeepSeek-V4 passes to attention: the sliding K/V cache
+    causal mask (seq×seq) concatenated with the compressed-cache block bias (seq×n_comp).
+
+    Mirrors ``DeepseekV4HCACompressor.forward`` (block_bias) + the local sliding window:
+    query t attends key j iff ``0 ≤ t-j < m`` (recent window), and compressed entry w iff
+    ``w < (t+1)//m`` (the window is closed/ready). Returns ``(grid, split_col, n_comp)`` where
+    columns ``[0:seq)`` are the sliding cache and ``[seq:seq+n_comp)`` the compressed cache."""
+    n_comp = max(seq // m, 0)
+    grid = []
+    for t in range(seq):
+        left = [1 if (j <= t and t - j < m) else 0 for j in range(seq)]
+        thr = (t + 1) // m
+        right = [1 if w < thr else 0 for w in range(n_comp)]
+        grid.append(left + right)
+    return grid, seq, n_comp
+
+
+def image_text_mask(n_img: int = 6, n_text: int = 8, image_bidirectional: bool = True):
+    """The cross-modal attention mask for a VLM: ``n_img`` image (prefix) tokens followed by
+    ``n_text`` text tokens. Text attends to ALL image tokens + causal text; image tokens are
+    bidirectional among themselves (prefix-LM, e.g. PaliGemma) or causal if not. Returns
+    ``(grid, split)`` with ``split = n_img`` (the image|text divider)."""
+    n = n_img + n_text
+    grid = []
+    for i in range(n):
+        row = []
+        for j in range(n):
+            if i < n_img:  # image query
+                attend = (j < n_img) if image_bidirectional else (j <= i)
+            else:  # text query: all image (prefix) + causal text
+                attend = (j < n_img) or (n_img <= j <= i)
+            row.append(1 if attend else 0)
+        grid.append(row)
+    return grid, n_img
+
+
+def attention_pattern_grids(tcfg, layer_types: list[str], seq: int = 24) -> dict[str, list[list[int]]]:
+    """Return ``{layer_type: grid}`` (schematic). grid[i][j]==1 ⇒ query i attends to key j.
+
+    Sliding/chunked are square (q×q). Compressed/heavily-compressed are RECTANGULAR (q ×
+    compressed-kv) to show that the KV cache is compressed -- the key point for DeepSeek-V4."""
+    types = list(dict.fromkeys(layer_types or ["full_attention"]))
+    rates = getattr(tcfg, "compress_rates", None) or {}
+    topk = getattr(tcfg, "index_topk", None)
+    out: dict[str, list[list[int]]] = {}
+    for lt in types:
+        n = lt.lower()
+        try:
+            if "compress" in n:  # compressed / heavily_compressed -> compressed KV cache
+                r = rates.get(lt, 4) if isinstance(rates, dict) else 4
+                out[lt] = _compressed_kv(seq, int(r), topk if "sparse" in n else None)
+            elif "chunk" in n:
+                out[lt] = _chunked(seq, _WIN)
+            elif "sliding" in n:
+                out[lt] = _sliding(seq, _WIN)
+            elif "linear" in n or "delta" in n or "mamba" in n or "recurrent" in n:
+                out[lt] = _causal(seq)
+            elif "bidirectional" in n:
+                out[lt] = _full(seq)
+            else:
+                out[lt] = _causal(seq)
+        except Exception:
+            continue
+    return out

arch_svg/modular.py

@@ -0,0 +1,322 @@
+"""The interesting part: extract *what a modular model changes* relative to its parent(s).
+
+We do not re-run the full ``utils/modular_model_converter.py`` (it exists to *generate*
+standalone modeling files). Instead we mirror its inheritance semantics with libcst to
+recover the diff payload directly:
+
+- a method/attr present in BOTH the modular class and the named parent class  -> overridden
+- present ONLY in the modular class                                            -> added
+- a deletion sentinel (``attr = AttributeError(...)`` or ``raise AttributeError``) -> deleted
+
+These are exactly the rules ``replace_class_node`` applies when it merges a child class
+onto its parent. The parent model a class inherits from is read from the modular file's
+``from ..<model>.modeling_<model> import <Class>`` imports -- the same import-following the
+converter does. A clean modular model yields a tiny ``ArchDiff``; a bloated one yields a
+large diff, which is the signal we want to surface.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass, field
+from functools import lru_cache
+
+import libcst as cst
+
+from .discover import models_root
+
+
+# ---------------------------------------------------------------------------- data model
+
+
+@dataclass
+class ClassChange:
+    """The change set for a single class defined in a modular file."""
+
+    name: str
+    relation: str  # "inherits" | "new"
+    parent: str | None = None
+    parent_model: str | None = None
+    overridden_methods: list[str] = field(default_factory=list)
+    added_methods: list[str] = field(default_factory=list)
+    deleted_methods: list[str] = field(default_factory=list)
+    overridden_attrs: list[str] = field(default_factory=list)
+    added_attrs: list[str] = field(default_factory=list)
+    deleted_attrs: list[str] = field(default_factory=list)
+
+    @property
+    def n_changes(self) -> int:
+        return (
+            len(self.overridden_methods)
+            + len(self.added_methods)
+            + len(self.deleted_methods)
+            + len(self.overridden_attrs)
+            + len(self.added_attrs)
+            + len(self.deleted_attrs)
+        )
+
+    @property
+    def is_trivial(self) -> bool:
+        """A pure ``class Foo(Bar): pass`` rename -- inherits everything, changes nothing."""
+        return self.relation == "inherits" and self.n_changes == 0
+
+
+@dataclass
+class ArchDiff:
+    model: str
+    is_modular: bool
+    parent_model: str | None
+    parent_models: list[str]
+    changes: list[ClassChange]
+    note: str | None = None
+
+    @property
+    def totals(self) -> dict[str, int]:
+        t = {"overridden": 0, "added": 0, "deleted": 0, "new_classes": 0, "trivial": 0}
+        for c in self.changes:
+            t["overridden"] += len(c.overridden_methods) + len(c.overridden_attrs)
+            t["added"] += len(c.added_methods) + len(c.added_attrs)
+            t["deleted"] += len(c.deleted_methods) + len(c.deleted_attrs)
+            if c.relation == "new":
+                t["new_classes"] += 1
+            if c.is_trivial:
+                t["trivial"] += 1
+        return t
+
+
+# ---------------------------------------------------------------------------- libcst bits
+
+
+def _dotted(node: cst.BaseExpression) -> str:
+    if isinstance(node, cst.Attribute):
+        return _dotted(node.value) + "." + node.attr.value
+    if isinstance(node, cst.Name):
+        return node.value
+    return ""
+
+
+def _base_name(base: cst.Arg) -> str:
+    v = base.value
+    if isinstance(v, cst.Name):
+        return v.value
+    if isinstance(v, cst.Attribute):
+        return v.attr.value
+    return _dotted(v)
+
+
+class _Collector(cst.CSTVisitor):
+    """Collect imports and top-level class definitions from a parsed module."""
+
+    def __init__(self) -> None:
+        # imported name -> source model dir (parsed from `..<model>.modeling_<model>`)
+        self.import_model: dict[str, str] = {}
+        self.classes: dict[str, cst.ClassDef] = {}
+        self._depth = 0
+
+    def visit_ImportFrom(self, node: cst.ImportFrom) -> None:
+        module = _dotted(node.module) if node.module is not None else ""
+        if ".modeling_" not in module:
+            return
+        parent_model = module.split(".modeling_")[0].rsplit(".", 1)[-1]
+        if isinstance(node.names, cst.ImportStar):
+            return
+        for alias in node.names:
+            self.import_model[alias.name.value] = parent_model
+
+    def visit_ClassDef(self, node: cst.ClassDef) -> bool:
+        if self._depth == 0:
+            self.classes[node.name.value] = node
+        self._depth += 1
+        return True
+
+    def leave_ClassDef(self, node: cst.ClassDef) -> None:
+        self._depth -= 1
+
+
+def _methods_and_attrs(klass: cst.ClassDef) -> tuple[dict[str, cst.FunctionDef], dict[str, cst.CSTNode]]:
+    methods: dict[str, cst.FunctionDef] = {}
+    attrs: dict[str, cst.CSTNode] = {}
+    for node in klass.body.body:
+        if isinstance(node, cst.FunctionDef):
+            methods[node.name.value] = node
+        elif isinstance(node, cst.SimpleStatementLine) and node.body:
+            stmt = node.body[0]
+            if isinstance(stmt, cst.Assign) and stmt.targets:
+                target = stmt.targets[0].target
+                if isinstance(target, cst.Name):
+                    attrs[target.value] = stmt
+            elif isinstance(stmt, cst.AnnAssign) and isinstance(stmt.target, cst.Name):
+                attrs[stmt.target.value] = stmt
+    return methods, attrs
+
+
+def _is_attr_deletion(stmt: cst.CSTNode) -> bool:
+    # `foo = AttributeError(...)` is the converter's sentinel for deleting an inherited attr.
+    if isinstance(stmt, cst.Assign) and isinstance(stmt.value, cst.Call):
+        return isinstance(stmt.value.func, cst.Name) and stmt.value.func.value == "AttributeError"
+    return False
+
+
+def _is_method_deletion(fn: cst.FunctionDef) -> bool:
+    # `def foo(self): raise AttributeError(...)` (or a bare `...`/del) signals removal.
+    body = fn.body.body if isinstance(fn.body, cst.IndentedBlock) else []
+    for node in body:
+        if isinstance(node, cst.SimpleStatementLine):
+            for s in node.body:
+                if isinstance(s, cst.Raise) and isinstance(s.exc, cst.Call):
+                    if isinstance(s.exc.func, cst.Name) and s.exc.func.value == "AttributeError":
+                        return True
+    return False
+
+
+@lru_cache(maxsize=512)
+def _parse_module_file(path: str) -> cst.Module | None:
+    try:
+        return cst.parse_module(open(path, encoding="utf-8").read())
+    except Exception:
+        return None
+
+
+@lru_cache(maxsize=512)
+def _parent_classes(parent_model: str) -> dict[str, cst.ClassDef]:
+    """Top-level classes defined in ``models/<parent_model>/modeling_<parent_model>.py``."""
+    path = os.path.join(models_root(), parent_model, f"modeling_{parent_model}.py")
+    if not os.path.exists(path):
+        # some models ship multiple modeling_*.py; scan them all
+        d = os.path.join(models_root(), parent_model)
+        if not os.path.isdir(d):
+            return {}
+        collector = _Collector()
+        for f in sorted(os.listdir(d)):
+            if f.startswith("modeling_") and f.endswith(".py"):
+                mod = _parse_module_file(os.path.join(d, f))
+                if mod is not None:
+                    mod.visit(collector)
+        return collector.classes
+    mod = _parse_module_file(path)
+    if mod is None:
+        return {}
+    collector = _Collector()
+    mod.visit(collector)
+    return collector.classes
+
+
+# ---------------------------------------------------------------------------- public API
+
+
+def _modular_path(model: str) -> str | None:
+    d = os.path.join(models_root(), model)
+    if not os.path.isdir(d):
+        return None
+    for f in sorted(os.listdir(d)):
+        if f.startswith("modular_") and f.endswith(".py"):
+            return os.path.join(d, f)
+    return None
+
+
+def resolve_parent(model: str) -> str | None:
+    """The dominant parent model of ``model`` (the one most classes inherit from), or None."""
+    d = diff(model)
+    return d.parent_model
+
+
+def diff(model: str) -> ArchDiff:
+    """Compute the modular diff for ``model``.
+
+    For a standalone model (no ``modular_*.py``), returns ``is_modular=False`` with a note.
+    """
+    path = _modular_path(model)
+    if path is None:
+        return ArchDiff(
+            model=model,
+            is_modular=False,
+            parent_model=None,
+            parent_models=[],
+            changes=[],
+            note="standalone model (no modular parent)",
+        )
+
+    mod = _parse_module_file(path)
+    if mod is None:
+        return ArchDiff(model, True, None, [], [], note="could not parse modular file")
+
+    collector = _Collector()
+    mod.visit(collector)
+
+    changes: list[ClassChange] = []
+    parent_vote: dict[str, int] = {}
+
+    for cls_name, node in collector.classes.items():
+        bases = [_base_name(b) for b in node.bases]
+        # find a base that maps to another transformers model
+        inherited_parent = next((b for b in bases if b in collector.import_model), None)
+
+        if inherited_parent is None:
+            # brand-new class (inherits nn.Module / PreTrainedConfig / a local class)
+            methods, attrs = _methods_and_attrs(node)
+            changes.append(
+                ClassChange(
+                    name=cls_name,
+                    relation="new",
+                    parent=bases[0] if bases else None,
+                    added_methods=sorted(methods),
+                    added_attrs=sorted(attrs),
+                )
+            )
+            continue
+
+        parent_model = collector.import_model[inherited_parent]
+        parent_vote[parent_model] = parent_vote.get(parent_model, 0) + 1
+        parent_cls = _parent_classes(parent_model).get(inherited_parent)
+
+        child_methods, child_attrs = _methods_and_attrs(node)
+        if parent_cls is None:
+            # could not locate parent source: treat all child members as overrides we cannot resolve
+            changes.append(
+                ClassChange(
+                    name=cls_name,
+                    relation="inherits",
+                    parent=inherited_parent,
+                    parent_model=parent_model,
+                    overridden_methods=sorted(child_methods),
+                    overridden_attrs=sorted(child_attrs),
+                )
+            )
+            continue
+
+        parent_methods, parent_attrs = _methods_and_attrs(parent_cls)
+        cc = ClassChange(name=cls_name, relation="inherits", parent=inherited_parent, parent_model=parent_model)
+        for name, fn in child_methods.items():
+            if _is_method_deletion(fn):
+                cc.deleted_methods.append(name)
+            elif name in parent_methods:
+                cc.overridden_methods.append(name)
+            else:
+                cc.added_methods.append(name)
+        for name, stmt in child_attrs.items():
+            if _is_attr_deletion(stmt):
+                cc.deleted_attrs.append(name)
+            elif name in parent_attrs:
+                cc.overridden_attrs.append(name)
+            else:
+                cc.added_attrs.append(name)
+        for lst in (
+            cc.overridden_methods,
+            cc.added_methods,
+            cc.deleted_methods,
+            cc.overridden_attrs,
+            cc.added_attrs,
+            cc.deleted_attrs,
+        ):
+            lst.sort()
+        changes.append(cc)
+
+    parent_model = max(parent_vote, key=parent_vote.get) if parent_vote else None
+    parent_models = sorted(parent_vote, key=lambda m: (-parent_vote[m], m))
+    return ArchDiff(
+        model=model,
+        is_modular=True,
+        parent_model=parent_model,
+        parent_models=parent_models,
+        changes=changes,
+    )

arch_svg/out/README.md

@@ -0,0 +1,22 @@
+---
+title: Transformers Model Architectures
+emoji: 📐
+colorFrom: indigo
+colorTo: blue
+sdk: static
+pinned: false
+license: apache-2.0
+---
+
+# 🤗 Transformers — Model Architecture Gallery
+
+Static, dependency-free architecture diagrams for every model in
+[`huggingface/transformers`](https://github.com/huggingface/transformers), generated by
+introspecting each model built on the `meta` device (no weights, no network).
+
+- **full** view — the decoder/encoder block decomposed top-to-bottom (attention or SSM mixer,
+  Q/K/V/O projections with real dims, norms, MLP / Sparse MoE, residual stream, mask grids).
+- **diff** view — for models with a `modular_*.py`, what each model overrides / adds / removes
+  versus the parent it inherits from.
+
+Open `index.html` for the browsable gallery.

arch_svg/render.py

@@ -0,0 +1,389 @@
+"""Positioned boxes + arrows -> SVG string.
+
+Theming is a one-line swap: all colors are CSS variables on ``:root``; flipping to the
+``.dark`` block (or honoring ``prefers-color-scheme``) reskins the whole gallery. Output is
+deterministic given a ``Diagram``.
+"""
+
+from __future__ import annotations
+
+from html import escape
+
+from .engine import fit_text
+from .layout import Diagram
+
+
+# palette: each component kind gets a stable fill/stroke variable pair
+_STYLE = """
+:root {
+  --bg: #ffffff; --fg: #1b1f24; --muted: #6b7280; --panel: #f6f8fa; --grid: #e5e7eb;
+  --embed: #dbeafe; --embed-s: #3b82f6;
+  --attn: #cffafe; --attn-s: #06b6d4;
+  --mamba: #dcfce7; --mamba-s: #22c55e;
+  --linattn: #fce7f3; --linattn-s: #ec4899;
+  --recur: #ede9fe; --recur-s: #8b5cf6;
+  --moe: #ffedd5; --moe-s: #f97316;
+  --mlp: #ede9fe; --mlp-s: #8b5cf6;
+  --norm: #e5e7eb; --norm-s: #9ca3af;
+  --head: #fee2e2; --head-s: #ef4444;
+  --config: #f1f5f9; --config-s: #64748b;
+  --rope: #fef9c3; --rope-s: #eab308;
+  --layer: #f8fafc; --layer-s: #cbd5e1;
+  --io: #f1f5f9; --io-s: #94a3b8;
+  --soft: #fae8ff; --soft-s: #c026d3;
+  --add: #ffffff; --add-s: #475569;
+  --block-s: #94a3b8;
+  --residual: #f59e0b;
+  --added: #16a34a; --over: #d97706; --deleted: #dc2626;
+  --lt-full: #06b6d4; --lt-sliding: #3b82f6; --lt-chunked: #8b5cf6;
+  --lt-compressed: #f97316; --lt-heavy: #dc2626; --lt-linear: #ec4899; --lt-mamba: #22c55e;
+  --cell-on: #0ea5e9; --cell-off: #e5e7eb;
+  --vision: #dcfce7; --vision-s: #16a34a; --audio: #fae8ff; --audio-s: #c026d3;
+  --proj: #fef3c7; --proj-s: #d97706; --xattn: #db2777;
+  --conv: #d1fae5; --conv-s: #10b981; --act: #ecfccb; --act-s: #65a30d;
+  --pool: #e0f2fe; --pool-s: #0284c7; --quant: #fae8ff; --quant-s: #c026d3;
+}
+@media (prefers-color-scheme: dark) {
+  :root {
+    --bg: #0d1117; --fg: #e6edf3; --muted: #8b949e; --panel: #161b22; --grid: #30363d;
+    --embed: #172554; --attn: #083344; --mamba: #052e16; --linattn: #500724;
+    --recur: #2e1065; --moe: #431407; --mlp: #2e1065; --norm: #21262d; --head: #450a0a;
+    --config: #1e293b; --rope: #422006; --layer: #161b22;
+    --conv: #022c22; --act: #1a2e05; --pool: #082f49; --quant: #3b0764; --proj: #422006;
+    --io: #1e293b; --soft: #3b0764; --add: #0d1117; --cell-off: #21262d;
+  }
+}
+.bg { fill: var(--bg); }
+text { font-family: ui-sans-serif, -apple-system, "Segoe UI", Roboto, sans-serif; fill: var(--fg); }
+.title { font-size: 22px; font-weight: 700; }
+.subtitle { font-size: 13px; fill: var(--muted); }
+.box-label { font-size: 14px; font-weight: 600; }
+.box-label.sm { font-size: 12.5px; }
+.box-sub { font-size: 11px; fill: var(--muted); }
+.glyph { font-size: 18px; font-weight: 700; fill: var(--add-s); }
+.badge { font-size: 12px; font-weight: 700; fill: var(--fg); }
+.facts-k { font-size: 11.5px; fill: var(--muted); }
+.facts-v { font-size: 11.5px; font-weight: 600; }
+.legend-t { font-size: 11.5px; fill: var(--fg); }
+.panel { fill: var(--panel); stroke: var(--grid); }
+rect.b { rx: 9; stroke-width: 1.6; }
+.c-embed { fill: var(--embed); stroke: var(--embed-s); }
+.c-attn  { fill: var(--attn);  stroke: var(--attn-s); }
+.c-mamba { fill: var(--mamba); stroke: var(--mamba-s); }
+.c-linattn { fill: var(--linattn); stroke: var(--linattn-s); }
+.c-recur { fill: var(--recur); stroke: var(--recur-s); }
+.c-moe   { fill: var(--moe);   stroke: var(--moe-s); }
+.c-mlp   { fill: var(--mlp);   stroke: var(--mlp-s); }
+.c-norm  { fill: var(--norm);  stroke: var(--norm-s); }
+.c-head  { fill: var(--head);  stroke: var(--head-s); }
+.c-config{ fill: var(--config);stroke: var(--config-s); }
+.c-rope  { fill: var(--rope);  stroke: var(--rope-s); }
+.c-proj  { fill: var(--proj);  stroke: var(--proj-s); }
+.c-conv  { fill: var(--conv);  stroke: var(--conv-s); }
+.c-act   { fill: var(--act);   stroke: var(--act-s); }
+.c-pool  { fill: var(--pool);  stroke: var(--pool-s); }
+.c-quant { fill: var(--quant); stroke: var(--quant-s); }
+.c-layer { fill: var(--layer); stroke: var(--layer-s); }
+.c-io    { fill: var(--io);    stroke: var(--io-s); }
+.c-soft  { fill: var(--soft);  stroke: var(--soft-s); }
+.c-add   { fill: var(--add);   stroke: var(--add-s); }
+.c-block { fill: none; stroke: var(--block-s); stroke-width: 1.6; stroke-dasharray: 7 5; }
+.c-lt-full { fill: var(--lt-full); stroke: var(--lt-full); }
+.c-lt-sliding { fill: var(--lt-sliding); stroke: var(--lt-sliding); }
+.c-lt-chunked { fill: var(--lt-chunked); stroke: var(--lt-chunked); }
+.c-lt-compressed { fill: var(--lt-compressed); stroke: var(--lt-compressed); }
+.c-lt-heavy { fill: var(--lt-heavy); stroke: var(--lt-heavy); }
+.c-lt-linear { fill: var(--lt-linear); stroke: var(--lt-linear); }
+.c-lt-mamba { fill: var(--lt-mamba); stroke: var(--lt-mamba); }
+.cell-on { fill: var(--cell-on); }
+.cell-off { fill: var(--cell-off); }
+.grid-frame { fill: none; stroke: var(--grid); stroke-width: 1; }
+.mask-bg { fill: var(--cell-off); }
+.mask-on { fill: #22c55e; }
+.mask-div { stroke: var(--fg); stroke-width: 1.5; stroke-dasharray: 3 2; }
+.c-vision { fill: var(--vision); stroke: var(--vision-s); }
+.c-audio  { fill: var(--audio);  stroke: var(--audio-s); }
+.c-proj   { fill: var(--proj);   stroke: var(--proj-s); }
+.c-sub    { fill: var(--bg); stroke: var(--block-s); stroke-width: 1.2; }
+.sec-h { font-size: 12px; font-weight: 700; }
+.sec-hbar { fill: var(--bg); opacity: 0.82; }
+.residual.xattn { stroke: var(--xattn); stroke-width: 2.4; }
+.ghost { opacity: 0.32; stroke-dasharray: 4 3; }
+.ch-added rect.b, rect.b.ch-added { stroke: var(--added); stroke-width: 3.2; }
+.ch-over rect.b, rect.b.ch-over { stroke: var(--over); stroke-width: 3.2; }
+.ch-deleted rect.b, rect.b.ch-deleted { stroke: var(--deleted); stroke-width: 3.2; }
+.edge { stroke: var(--grid); stroke-width: 2; }
+.flow { stroke: var(--grid); stroke-width: 2; fill: none; }
+.residual { stroke: var(--residual); stroke-width: 2; fill: none; }
+.rope { stroke: var(--rope-s); stroke-width: 2.2; fill: none; }
+.xattn { stroke: var(--xattn); stroke-width: 2.4; fill: none; }
+.cell-idx { font-size: 9px; fill: #ffffff; font-weight: 600; }
+.sky { fill: #bae6fd; } .sun { fill: #fde047; } .hill { fill: #4ade80; }
+"""
+
+_CHANGE_SWATCH = {"added": "var(--added)", "overridden": "var(--over)", "deleted": "var(--deleted)"}
+_LEGEND_SWATCH = {
+    "c-attn": "var(--attn-s)",
+    "c-moe": "var(--moe-s)",
+    "c-mlp": "var(--mlp-s)",
+    "c-mamba": "var(--mamba-s)",
+    "c-norm": "var(--norm-s)",
+    "c-embed": "var(--embed-s)",
+    "c-head": "var(--head-s)",
+    "c-rope": "var(--rope-s)",
+    "c-vision": "var(--vision-s)",
+    "c-audio": "var(--audio-s)",
+    "c-proj": "var(--proj-s)",
+    "xattn": "var(--xattn)",
+    "c-lt-full": "var(--lt-full)",
+    "c-lt-sliding": "var(--lt-sliding)",
+    "c-lt-chunked": "var(--lt-chunked)",
+    "c-lt-compressed": "var(--lt-compressed)",
+    "c-lt-heavy": "var(--lt-heavy)",
+    "c-lt-linear": "var(--lt-linear)",
+    "c-lt-mamba": "var(--lt-mamba)",
+    "residual": "var(--residual)",
+    "ch-added": "var(--added)",
+    "ch-over": "var(--over)",
+    "ch-deleted": "var(--deleted)",
+    "ghost": "var(--grid)",
+}
+
+
+def _box_svg(b) -> str:
+    change_cls = {"added": "ch-added", "overridden": "ch-over", "deleted": "ch-deleted"}.get(b.change, "")
+    group_cls = "ghost" if b.ghost else ""
+    parts = [f'<g class="{group_cls}">']
+    if b.title:
+        parts.append(f"<title>{escape(str(b.title))}</title>")
+    cx = b.x + b.w / 2
+
+    if b.shape == "cell":  # a single layer-schedule tick with its index
+        parts.append(f'<rect class="{b.cls}" x="{b.x}" y="{b.y}" width="{b.w}" height="{b.h}" rx="2"/>')
+        if b.label and b.h >= 11:
+            parts.append(
+                f'<text class="cell-idx" x="{b.x + b.w / 2}" y="{b.y + b.h / 2 + 3.5}" text-anchor="middle">{escape(b.label)}</text>'
+            )
+        parts.append("</g>")
+        return "".join(parts)
+
+    if b.shape == "grid" and b.grid:
+        rows = len(b.grid)
+        cols = len(b.grid[0]) if rows else 0
+        cell = b.w / max(cols, 1)  # square cells; width fixed, height follows rows
+        parts.append(
+            f'<text class="box-sub" x="{b.x}" y="{b.y - 6}">{escape(fit_text(b.label, max(b.w, 130), 11))}</text>'
+        )
+        parts.append(
+            f'<rect class="mask-bg" x="{b.x}" y="{b.y}" width="{cell * cols:.1f}" height="{cell * rows:.1f}"/>'
+        )
+        for i, row in enumerate(b.grid):
+            for j, c in enumerate(row):
+                if c:
+                    parts.append(
+                        f'<rect class="mask-on" x="{b.x + j * cell:.1f}" y="{b.y + i * cell:.1f}" '
+                        f'width="{cell:.1f}" height="{cell:.1f}"/>'
+                    )
+        if b.grid_split is not None and 0 < b.grid_split < cols:  # sliding | compressed divider
+            dx = b.x + b.grid_split * cell
+            parts.append(
+                f'<line class="mask-div" x1="{dx:.1f}" y1="{b.y}" x2="{dx:.1f}" y2="{b.y + cell * rows:.1f}"/>'
+            )
+        parts.append(
+            f'<rect class="grid-frame" x="{b.x}" y="{b.y}" width="{cell * cols:.1f}" height="{cell * rows:.1f}"/>'
+        )
+        parts.append("</g>")
+        return "".join(parts)
+
+    if b.shape == "circle":
+        r = b.h // 2
+        parts.append(f'<circle class="b {b.cls} {change_cls}" cx="{cx}" cy="{b.y + r}" r="{r}"/>')
+        if b.glyph:
+            parts.append(
+                f'<text class="glyph" x="{cx}" y="{b.y + r + 6}" text-anchor="middle">{escape(b.glyph)}</text>'
+            )
+        parts.append("</g>")
+        return "".join(parts)
+
+    if b.shape == "container":
+        parts.append(f'<rect class="{b.cls}" x="{b.x}" y="{b.y}" width="{b.w}" height="{b.h}" rx="14"/>')
+        parts.append(f'<text class="box-sub" x="{b.x + 12}" y="{b.y + 18}">{escape(b.label)}</text>')
+        if b.badge:
+            parts.append(
+                f'<text class="badge" x="{b.x + b.w - 12}" y="{b.y + 20}" text-anchor="end">{escape(b.badge)}</text>'
+            )
+        parts.append("</g>")
+        return "".join(parts)
+
+    if b.shape == "section":  # filled rounded panel with a header; inner chips drawn on top
+        change_cls = {"added": "ch-added", "overridden": "ch-over", "deleted": "ch-deleted"}.get(b.change, "")
+        parts.append(
+            f'<rect class="b {b.cls} {change_cls}" x="{b.x}" y="{b.y}" width="{b.w}" height="{b.h}" rx="11"/>'
+        )
+        # white header strip so the class name stays readable on saturated section fills
+        parts.append(f'<rect class="sec-hbar" x="{b.x + 5}" y="{b.y + 4}" width="{b.w - 10}" height="18" rx="5"/>')
+        hdr = fit_text(b.label, b.w - 24 - (40 if b.badge else 0), 12)
+        parts.append(f'<text class="sec-h" x="{b.x + 12}" y="{b.y + 17}">{escape(hdr)}</text>')
+        if b.badge:
+            parts.append(
+                f'<text class="badge" x="{b.x + b.w - 12}" y="{b.y + 17}" text-anchor="end">{escape(b.badge)}</text>'
+            )
+        parts.append("</g>")
+        return "".join(parts)
+
+    if b.shape == "image":  # tiny schematic example image (for pixel_values)
+        parts.append(f'<rect class="b {b.cls}" x="{b.x}" y="{b.y}" width="{b.w}" height="{b.h}" rx="6"/>')
+        ix, iy, iw, ih = b.x + 6, b.y + 6, b.w - 12, b.h - 12
+        parts.append(
+            f'<clipPath id="ic{b.x}{b.y}"><rect x="{ix}" y="{iy}" width="{iw}" height="{ih}" rx="3"/></clipPath>'
+        )
+        g = f'<g clip-path="url(#ic{b.x}{b.y})">'
+        g += f'<rect class="sky" x="{ix}" y="{iy}" width="{iw}" height="{ih}"/>'
+        g += f'<circle class="sun" cx="{ix + iw * 0.72:.0f}" cy="{iy + ih * 0.32:.0f}" r="{ih * 0.16:.0f}"/>'
+        g += f'<path class="hill" d="M{ix},{iy + ih} L{ix + iw * 0.35:.0f},{iy + ih * 0.5:.0f} L{ix + iw * 0.6:.0f},{iy + ih:.0f} Z"/>'
+        g += f'<path class="hill" d="M{ix + iw * 0.45:.0f},{iy + ih} L{ix + iw * 0.75:.0f},{iy + ih * 0.55:.0f} L{ix + iw:.0f},{iy + ih:.0f} Z"/>'
+        g += "</g>"
+        parts.append(g)
+        parts.append(f'<rect class="grid-frame" x="{ix}" y="{iy}" width="{iw}" height="{ih}" rx="3"/>')
+        if b.label:
+            parts.append(
+                f'<text class="box-sub" x="{cx}" y="{b.y + b.h + 12}" text-anchor="middle">{escape(b.label)}</text>'
+            )
+        parts.append("</g>")
+        return "".join(parts)
+
+    rx = b.h // 2 if b.shape == "io" else 9
+    rect_cls = f"b {b.cls} {change_cls}".strip()
+    parts.append(f'<rect class="{rect_cls}" x="{b.x}" y="{b.y}" width="{b.w}" height="{b.h}" rx="{rx}"/>')
+    label_cls = "box-label sm" if b.small else "box-label"
+    lpx = 12.5 if b.small else 14
+    inner = b.w - 14  # text must fit inside the box with a little padding
+    # vertically centre the label + sublabels block within the box
+    subs = b.sublabels or []
+    n_sub = len(subs)
+    if n_sub == 0:
+        ty = b.y + b.h / 2 + 4.5
+    else:
+        block = 13 + 15 * n_sub
+        ty = b.y + (b.h - block) / 2 + 12
+    lab = fit_text(b.label, inner, lpx)
+    parts.append(f'<text class="{label_cls}" x="{cx}" y="{ty:.1f}" text-anchor="middle">{escape(lab)}</text>')
+    sy = ty + 15
+    for s in subs:
+        parts.append(
+            f'<text class="box-sub" x="{cx}" y="{sy:.1f}" text-anchor="middle">{escape(fit_text(s, inner, 11))}</text>'
+        )
+        sy += 15
+    if b.badge:
+        parts.append(
+            f'<text class="badge" x="{b.x + b.w - 12}" y="{b.y + 20}" text-anchor="end">{escape(b.badge)}</text>'
+        )
+    if b.change:
+        sw = _CHANGE_SWATCH.get(b.change, "")
+        parts.append(f'<circle cx="{b.x + 12}" cy="{b.y + 12}" r="5" fill="{sw}"/>')
+    parts.append("</g>")
+    return "".join(parts)
+
+
+def _arrow_svg(a) -> str:
+    pts = " ".join(f"{x},{y}" for x, y in a.points)
+    dash = ' stroke-dasharray="6 4"' if a.dashed else ""
+    return f'<polyline class="{a.cls}" points="{pts}"{dash} marker-end="url(#ah-{a.cls})"/>'
+
+
+def _marker(name: str, color: str) -> str:
+    return (
+        f'<marker id="ah-{name}" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="7" markerHeight="7" '
+        f'orient="auto-start-reverse"><path d="M0,0 L10,5 L0,10 z" fill="{color}"/></marker>'
+    )
+
+
+def render(diagram: Diagram) -> str:
+    d = diagram
+    panel_w = 264
+    panel_x = d.width - panel_w - 24
+    parts = [
+        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {d.width} {d.height}" '
+        f'width="{d.width}" height="{d.height}" font-size="14">',
+        f"<style>{_STYLE}</style>",
+        f"<defs>{_marker('flow', 'var(--grid)')}{_marker('residual', 'var(--residual)')}"
+        f"{_marker('rope', 'var(--rope-s)')}{_marker('xattn', 'var(--xattn)')}</defs>",
+        f'<rect class="bg" x="0" y="0" width="{d.width}" height="{d.height}"/>',
+        f'<text class="title" x="24" y="34">{escape(d.title)}</text>',
+        f'<text class="subtitle" x="24" y="54">{escape(d.subtitle)}</text>',
+    ]
+
+    # connecting spine behind the central column only (use the modal box-center x)
+    if d.spine:
+        col = [b for b in d.boxes if b.shape not in ("container", "cell", "grid")]
+        if col:
+            from collections import Counter
+
+            centers = Counter(round(b.x + b.w / 2) for b in col)
+            xs = centers.most_common(1)[0][0]
+            on_axis = [b for b in col if abs((b.x + b.w / 2) - xs) < 6]
+            top = min(b.y for b in on_axis)
+            bot = max(b.y + b.h for b in on_axis)
+            parts.append(f'<line class="edge" x1="{xs}" y1="{top}" x2="{xs}" y2="{bot}"/>')
+
+    # containers + sections first (behind), then arrows, then boxes on top.
+    # sections are drawn LARGEST-first so a nested (inner) section paints on top of its
+    # parent instead of being hidden by the parent's opaque fill.
+    for b in sorted((b for b in d.boxes if b.shape == "container"), key=lambda b: -b.w * b.h):
+        parts.append(_box_svg(b))
+    for b in sorted((b for b in d.boxes if b.shape == "section"), key=lambda b: -b.w * b.h):
+        parts.append(_box_svg(b))
+    for a in d.arrows:
+        parts.append(_arrow_svg(a))
+    for b in d.boxes:
+        if b.shape not in ("container", "section"):
+            parts.append(_box_svg(b))
+
+    # side panel: facts
+    py = 88
+    parts.append(
+        f'<rect class="panel" x="{panel_x}" y="{py}" width="{panel_w}" height="{len(d.facts) * 22 + 20}" rx="8"/>'
+    )
+    fy = py + 24
+    for k, v in d.facts:
+        parts.append(f'<text class="facts-k" x="{panel_x + 14}" y="{fy}">{escape(k)}</text>')
+        parts.append(
+            f'<text class="facts-v" x="{panel_x + panel_w - 14}" y="{fy}" text-anchor="end">{escape(str(v))}</text>'
+        )
+        fy += 22
+    fy += 16
+
+    # side panel: legend
+    parts.append(f'<text class="legend-t" x="{panel_x + 2}" y="{fy}" font-weight="700">legend</text>')
+    fy += 12
+    _seen_leg = set()
+    legend = [(c, l) for c, l in d.legend if not ((c, l) in _seen_leg or _seen_leg.add((c, l)))]
+    for cls, label in legend:
+        sw = _LEGEND_SWATCH.get(cls, "var(--grid)")
+        is_change = cls.startswith("ch-") or cls in ("ghost", "residual")
+        if is_change:
+            parts.append(
+                f'<rect x="{panel_x + 4}" y="{fy - 1}" width="16" height="12" rx="2" fill="none" stroke="{sw}" stroke-width="3"/>'
+            )
+        else:
+            parts.append(f'<rect x="{panel_x + 4}" y="{fy - 1}" width="16" height="12" rx="2" fill="{sw}"/>')
+        parts.append(f'<text class="legend-t" x="{panel_x + 28}" y="{fy + 9}">{escape(label)}</text>')
+        fy += 18
+
+    # side panel: per-class changes (diff mode)
+    if d.changes:
+        fy += 14
+        parts.append(f'<text class="legend-t" x="{panel_x + 2}" y="{fy}" font-weight="700">changes by class</text>')
+        fy += 16
+        for ch, name, detail in d.changes:
+            sw = _CHANGE_SWATCH.get(ch, "var(--grid)")
+            parts.append(f'<circle cx="{panel_x + 8}" cy="{fy - 3}" r="4" fill="{sw}"/>')
+            parts.append(f'<text class="facts-v" x="{panel_x + 20}" y="{fy}">{escape(name)}</text>')
+            fy += 13
+            if detail:
+                d_short = detail if len(detail) <= 46 else detail[:44] + "…"
+                parts.append(f'<text class="box-sub" x="{panel_x + 20}" y="{fy}">{escape(d_short)}</text>')
+                fy += 15
+
+    parts.append("</svg>")
+    return "\n".join(parts)