inline-core 1.2.66__tar.gz → 1.2.68__tar.gz

{inline_core-1.2.66 → inline_core-1.2.68}/CLAUDE.md

@@ -163,6 +163,17 @@ between nodes and are never takes.
   → sequential offload → wont-fit. NF4 is what makes a 32B model viable on a 24 GB card. Note int8
   forces bf16 (torchao's weight-only int8 silently no-ops under fp16) while NF4 does not, so a Turing
   card keeps its fp16 tensor cores under NF4.
+- **A pre-reduced checkpoint must not be re-reduced, structurally or numerically.** MiniMax H3's
+  `pruned` builds ship the AdaLN branch already factorised to rank 8 and drop the timestep path
+  entirely, so re-running our factorisation multiplies a `[96768, 8]` projection by a full-width
+  basis. `minimaxh3/pipeline.py` turns `factorise_adaln` off for those, the same way the rule below
+  turns quantization off for a prequantized file. Both are the same rule: the source is already in
+  the target form.
+- **Size a checkpoint by what it becomes, not by what it weighs.** A pruned file has already lost
+  its AdaLN branch and an fp8 file stores half the bytes it will occupy once dequantised, so scaling
+  the on-disk number under-sizes both, by up to 3x. `minimaxh3.requirements.resident_bytes` counts
+  from the header. Under-sizing is the dangerous direction: the fit ladder then promises a machine
+  that dies to a host-RAM OOM kill instead of raising.
 - **Prequantized checkpoints must not be re-quantized.** A checkpoint that ships already quantized
   (`flux2/variants.is_prequantized`) has an on-disk size that already _is_ its resident size, so the
   ladder's assumption that quantization halves it does not hold, and handing diffusers a second,
@@ -230,6 +241,16 @@ real codec that moves tensors lives with the model runner.
   Don't scatter it.
 - **Bring-your-own models.** Nothing is downloaded by the engine. The catalog scans; the user places
   files. A model picker is a `SELECT` param with `options_from="<category>"`.
+- **Adapter strength is not a quality metric, and a threshold on it is a false-positive machine.**
+  Measured against real bases, published LoRAs that work well span `|B@A| / |W|` from 0.017%
+  (a style LoRA) to 1.2% (a restoration LoRA), so "this adapter looks weak" is not a finding. What
+  predicts a LoRA doing nothing is whether its per-weight change clears one quantization step. Warn
+  on that, and only when the base is actually quantized.
+- **Patching a diffusers object may patch a copy, and it will not tell you.**
+  `ModularPipeline.blocks` is a property returning `deepcopy(self._blocks)`, so hooking the block
+  graph through it installs cleanly onto a throwaway and reports success. That is how H3's denoise
+  ran with no per-step progress while the hook said it was attached. Reach for the backing
+  attribute, and prove a hook fires against the real object rather than trusting a return value.
 - **Verify image models by rendering.** The FLUX.2 work shipped five bugs past a green test suite,
   and every one produced a _wrong image rather than an error_: a mis-keyed checkpoint, a
   vision-language encoder loaded in place of a text one, an unnormalized latent, and a control context

{inline_core-1.2.66 → inline_core-1.2.68}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: inline-core
-Version: 1.2.66
+Version: 1.2.68
 Summary: The generation engine behind Inline Studio.
 License-Expression: GPL-3.0-or-later
 Requires-Python: >=3.11

{inline_core-1.2.66 → inline_core-1.2.68}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 # PyPI name; the import package is `inline_core` (src/inline_core).
 name = "inline-core"
-version = "1.2.66"
+version = "1.2.68"
 description = "The generation engine behind Inline Studio."
 readme = "README.md"
 license = "GPL-3.0-or-later"

{inline_core-1.2.66 → inline_core-1.2.68}/src/inline_core/config.py

@@ -7,12 +7,33 @@ from pathlib import Path
 
 
 def models_dir() -> Path:
-    """The models root scanned on start (category subfolders inside). `INLINE_MODELS_DIR`, else
-    `./models`. Users drop their own weight files here; nothing is downloaded."""
+    """The **writable** models root: where downloads and trained LoRAs land. `INLINE_MODELS_DIR`,
+    else `./models`. Users drop their own weight files here; nothing is downloaded."""
     env = os.environ.get("INLINE_MODELS_DIR")
     return Path(env).expanduser() if env else Path("models")
 
 
+def models_dirs() -> list[Path]:
+    """Every root scanned for installed weights, writable one first.
+
+    Pointing `INLINE_MODELS_DIR` (or `--models-dir`) somewhere else used to hide `./models`
+    entirely, so a user with weights in both places could only ever see one set. The extra roots
+    are read-only: only `models_dir()` is ever written to.
+    """
+    roots = [models_dir()]
+    default = Path("models")
+    for extra in (default, *_extra_models_dirs()):
+        if extra not in roots and extra.is_dir():
+            roots.append(extra)
+    return roots
+
+
+def _extra_models_dirs() -> list[Path]:
+    """Additional read-only roots from `INLINE_EXTRA_MODELS_DIRS` (os.pathsep-separated)."""
+    raw = os.environ.get("INLINE_EXTRA_MODELS_DIRS", "")
+    return [Path(p).expanduser() for p in raw.split(os.pathsep) if p.strip()]
+
+
 def data_dir() -> Path:
     """Engine-owned working data (the run DB, takes). `INLINE_DATA_DIR`, else `./.inline`."""
     env = os.environ.get("INLINE_DATA_DIR")

{inline_core-1.2.66 → inline_core-1.2.68}/src/inline_core/models/catalog.py

@@ -13,9 +13,11 @@ from __future__ import annotations
 
 import hashlib
 import json
+from collections.abc import Sequence
 from pathlib import Path
+from typing import Any
 
-from ..config import models_dir
+from ..config import models_dirs
 
 # Category subfolders scanned under the models root. These are the keys a param's `options_from`
 # may reference (see graph/primitives.py); ensure_dirs() creates them so drop-in is obvious.
@@ -29,6 +31,9 @@ CATEGORIES: tuple[str, ...] = (
     "controlnet",
     "upscale_models",
     "embeddings",
+    # Written by the preprocess runner rather than the user, but it holds real weights and the
+    # panel is meant to show everything on disk.
+    "annotators",
 )
 
 # Extensions we treat as model weights. A folder counts as a model if it contains one of these.
@@ -45,34 +50,111 @@ def _folder_has_weight(path: Path) -> bool:
     return any(_is_weight(child) for child in path.rglob("*"))
 
 
+def _file_node(path: Path) -> dict[str, Any]:
+    stat = path.stat()
+    return {
+        "name": path.name,
+        "path": str(path),
+        "kind": "file",
+        "sizeBytes": stat.st_size,
+        "mtime": int(stat.st_mtime * 1000),
+    }
+
+
+def _dir_node(path: Path, name: str) -> dict[str, Any] | None:
+    """A directory as a tree node, or None when it holds nothing worth showing."""
+    if not path.is_dir():
+        return None
+    children: list[dict[str, Any]] = []
+    try:
+        entries = sorted(path.iterdir(), key=lambda p: p.name.lower())
+    except OSError:
+        return None
+    for entry in entries:
+        # Caches like loras/.cache and text_encoders/.snap are noise in a panel about weights.
+        if entry.name.startswith("."):
+            continue
+        if _is_weight(entry):
+            children.append(_file_node(entry))
+        elif entry.is_dir():
+            child = _dir_node(entry, entry.name)
+            if child is not None:
+                children.append(child)
+    if not children:
+        return None
+    return {
+        "name": name,
+        "path": str(path),
+        "kind": "dir",
+        "children": children,
+        "fileCount": sum(1 for c in children if c["kind"] == "file"),
+    }
+
+
 class ModelCatalog:
-    """Scans the models root and answers "what's installed" per category.
+    """Scans the models roots and answers "what's installed" per category.
 
     Cheap to construct; nothing touches disk until ``ensure_dirs`` or ``rescan``/``scan``. Results
-    are cached so ``list`` and ``fingerprint`` are hits between scans.
+    are cached so ``list`` and ``fingerprint`` are hits between scans. Several roots may be scanned
+    (see ``config.models_dirs``); only the first is ever written to.
     """
 
-    def __init__(self, root: Path) -> None:
-        self._root = Path(root)
+    def __init__(self, root: Path | Sequence[Path]) -> None:
+        roots = [Path(root)] if isinstance(root, str | Path) else [Path(r) for r in root]
+        self._roots = roots or [Path("models")]
         self._entries: dict[str, list[str]] = {category: [] for category in CATEGORIES}
 
     @property
     def root(self) -> Path:
-        return self._root
+        """The writable root. Downloads and trained LoRAs land here."""
+        return self._roots[0]
+
+    @property
+    def roots(self) -> list[Path]:
+        return list(self._roots)
 
     def ensure_dirs(self) -> None:
-        """Create the root and every category subfolder, so users have somewhere to drop weights."""
+        """Create the writable root and its category subfolders, so drop-in has somewhere to go."""
         for category in CATEGORIES:
-            (self._root / category).mkdir(parents=True, exist_ok=True)
+            (self.root / category).mkdir(parents=True, exist_ok=True)
 
     def rescan(self) -> dict[str, list[str]]:
-        """Re-read every category from disk, cache the result, and return it."""
+        """Re-read every category from every root, cache the result, and return it."""
         entries: dict[str, list[str]] = {}
         for category in CATEGORIES:
-            entries[category] = self._scan_category(self._root / category)
+            names: list[str] = []
+            for root in self._roots:
+                for name in self._scan_category(root / category):
+                    if name not in names:
+                        names.append(name)
+            entries[category] = sorted(names)
         self._entries = entries
         return entries
 
+    def tree(self) -> list[dict[str, Any]]:
+        """Every root as a nested, read-only listing for the Models panel.
+
+        Separate from ``rescan`` because the select only needs flat names, while the panel wants
+        nesting, sizes and mtimes; keeping them apart stops the hot path paying for the detail.
+        """
+        out: list[dict[str, Any]] = []
+        for index, root in enumerate(self._roots):
+            categories = [
+                node
+                for category in CATEGORIES
+                if (node := _dir_node(root / category, category)) is not None
+            ]
+            out.append(
+                {
+                    "path": str(root),
+                    "label": root.name or str(root),
+                    "writable": index == 0,
+                    "exists": root.is_dir(),
+                    "categories": categories,
+                }
+            )
+        return out
+
     # app.py calls scan() in the lifespan; rescan() is the same work exposed for tests/callers that
     # want the mapping back. Keep both so neither call site has to know about the other.
     def scan(self) -> dict[str, list[str]]:
@@ -92,6 +174,8 @@ class ModelCatalog:
             return []
         names: list[str] = []
         for entry in directory.iterdir():
+            if entry.name.startswith("."):
+                continue
             if _is_weight(entry):
                 names.append(entry.name)
             elif entry.is_dir() and _folder_has_weight(entry):
@@ -112,8 +196,11 @@ def resolve_picked(category: str, chosen: object) -> Path | None:
     name = str(chosen or "").strip()
     if not name:
         return None
-    root = models_dir() / category
-    for candidate in (root / name, root / Path(name).name, Path(name)):
-        if candidate.exists():
-            return candidate
-    return None
+    # Every root, writable one first, so a pick resolves wherever the user actually keeps it.
+    for models_root in models_dirs():
+        root = models_root / category
+        for candidate in (root / name, root / Path(name).name):
+            if candidate.exists():
+                return candidate
+    bare = Path(name)
+    return bare if bare.exists() else None

{inline_core-1.2.66 → inline_core-1.2.68}/src/inline_core/models/keymap.py

@@ -262,7 +262,7 @@ def transform(
         raise ComponentError(f"{key} has {tensor.shape[0]} rows, not divisible into {parts} parts.")
     if verify_layout:
         assert_layout(tensor, action, key=key)
-    source = _deinterleave(tensor, parts, action.head_dim) if (
+    source = deinterleave_rows(tensor, parts, action.head_dim) if (
         action.layout is RowLayout.INTERLEAVED
     ) else tensor
     block = source.shape[0] // parts
@@ -270,16 +270,31 @@ def transform(
         yield target, source[index * block : (index + 1) * block]
 
 
-def _deinterleave(tensor: Any, parts: int, head_dim: int) -> Any:
-    """``[p0_h0; p1_h0; p2_h0][p0_h1; …]`` to ``[p0_all; p1_all; p2_all]``.
-
-    ``transpose`` is not the same call in torch and numpy - torch swaps two axes, numpy wants a full
-    permutation - so the swap is spelled per backend rather than duck-typed.
-    """
+def deinterleave_rows(tensor: Any, parts: int, head_dim: int) -> Any:
+    """``[p0_h0; p1_h0; p2_h0][p0_h1; …]`` to ``[p0_all; p1_all; p2_all]``."""
     if head_dim < 1:
         raise ComponentError("De-interleaving needs the head dimension the parts are grouped by.")
     heads = tensor.shape[0] // (parts * head_dim)
-    reshaped = tensor.reshape(heads, parts, head_dim, *tensor.shape[1:])
+    return _swap01(tensor, (heads, parts, head_dim))
+
+
+def interleave_rows(tensor: Any, parts: int, head_dim: int) -> Any:
+    """``[p0_all; p1_all; p2_all]`` back to per-head groups: the inverse of ``deinterleave_rows``.
+
+    Needed to *write* a checkpoint or adapter in a publisher's interleaved layout, where the load
+    path only ever reads one."""
+    if head_dim < 1:
+        raise ComponentError("Interleaving needs the head dimension the parts are grouped by.")
+    heads = tensor.shape[0] // (parts * head_dim)
+    return _swap01(tensor, (parts, heads, head_dim))
+
+
+def _swap01(tensor: Any, shape: tuple[int, int, int]) -> Any:
+    """Reshape to ``shape`` plus the trailing dims, exchange the first two, flatten back.
+
+    ``transpose`` is not the same call in torch and numpy - torch swaps two axes, numpy wants a full
+    permutation - so the swap is spelled per backend rather than duck-typed."""
+    reshaped = tensor.reshape(*shape, *tensor.shape[1:])
     if _is_torch(tensor):
         moved = reshaped.transpose(0, 1).contiguous()
     else:

{inline_core-1.2.66 → inline_core-1.2.68}/src/inline_core/models/lora.py

@@ -29,6 +29,10 @@ if TYPE_CHECKING:
 #: while ostris' training adapter uses the reference names.
 Alias = Callable[[str], str | None]
 
+#: Rewrites a whole adapter before it is matched, for an arch whose checkpoint keys need more than a
+#: rename - MiniMax H3 ships attention fused, so three of our modules are one of theirs.
+Translate = Callable[[dict[str, Any]], dict[str, Any]]
+
 _DOWN = ("lora_down.weight", "lora_A.weight", "lora_A.default.weight")
 _UP = ("lora_up.weight", "lora_B.weight", "lora_B.default.weight")
 # Prefixes checkpoints put in front of the module path; stripped when matching against the model.
@@ -39,12 +43,22 @@ _PREFIXES = ("diffusion_model.", "transformer.", "lora_unet_", "lora_te_", "base
 LoraPlan = dict[str, list[tuple[Any, Any, float]]]
 
 
-def fuse_loras(model: Any, loras: tuple[LoraRef, ...], alias: Alias | None = None) -> None:
+def fuse_loras(
+    model: Any,
+    loras: tuple[LoraRef, ...],
+    alias: Alias | None = None,
+    translate: Translate | None = None,
+) -> None:
     """Merge each LoRA into ``model``'s weights in order. No-op for an empty stack."""
-    apply_plan(model, plan_loras(model, loras, alias))
+    apply_plan(model, plan_loras(model, loras, alias, translate))
 
 
-def plan_loras(model: Any, loras: tuple[LoraRef, ...], alias: Alias | None = None) -> LoraPlan:
+def plan_loras(
+    model: Any,
+    loras: tuple[LoraRef, ...],
+    alias: Alias | None = None,
+    translate: Translate | None = None,
+) -> LoraPlan:
     """Resolve every LoRA against ``model``'s module names, without touching any weights.
 
     Split from the fusing so a streaming loader can validate the whole stack **before** reading a
@@ -53,7 +67,7 @@ def plan_loras(model: Any, loras: tuple[LoraRef, ...], alias: Alias | None = Non
     plan: LoraPlan = {}
     names = _linear_module_names(model)
     for lora in loras:
-        _plan_one(plan, names, lora.file, lora.strength, alias)
+        _plan_one(plan, names, lora.file, lora.strength, alias, translate)
     return plan
 
 
@@ -74,7 +88,12 @@ def apply_plan(module: Any, plan: LoraPlan, prefix: str = "") -> None:
 
 
 def _plan_one(
-    plan: LoraPlan, names: dict[str, None], path: str, strength: float, alias: Alias | None
+    plan: LoraPlan,
+    names: dict[str, None],
+    path: str,
+    strength: float,
+    alias: Alias | None,
+    translate: Translate | None = None,
 ) -> None:
     from safetensors.torch import load_file
 
@@ -83,6 +102,8 @@ def _plan_one(
     except Exception as exc:  # noqa: BLE001
         raise ComponentError(f"Could not read LoRA {path!r}: {exc}") from exc
 
+    if translate is not None:
+        state = translate(state)
     pairs, alphas = _group(state)
     if not pairs:
         raise ComponentError(f"LoRA {path!r} contains no recognisable lora_down/lora_up pairs.")
@@ -104,27 +125,39 @@ def _plan_one(
         )
 
 
+#: Most fp32 delta held at once. The product is computed a slice of output rows at a time because
+#: the whole of it is enormous: one MiniMax H3 block's six Linears come to 1.5GB and the model to
+#: 80GB, which is host RAM during a staged load and took a 60GB box down three times.
+_DELTA_CHUNK_BYTES = 64 * 1024 * 1024
+
+
 def _add_delta(weight: Any, up: Any, down: Any, scale: float) -> None:
     """Fuse ``scale * (up @ down)`` into ``weight`` in place.
 
-    Computed on the weight's own device: a big LoRA (Krea 2's are ~260 modules on a 12.9B model)
-    materializes tens of GB of fp32 deltas, and doing that on the CPU costs ~20s of maths plus the
-    transfer where the GPU takes ~2s. Falls back to the CPU if the device runs out of memory, so a
-    tight card still fuses, just slowly."""
+    Computed on the weight's own device: doing it on the CPU costs ~20s of maths plus the transfer
+    where the GPU takes ~2s. Falls back to the CPU if the device runs out of memory, so a tight card
+    still fuses, just slowly."""
     import torch
 
     try:
-        weight.add_(_delta(up, down, weight, weight.device) * scale)
+        _accumulate(weight, up, down, scale, weight.device)
     except torch.cuda.OutOfMemoryError:
-        weight.add_((_delta(up, down, weight, "cpu") * scale).to(weight.device, weight.dtype))
+        _accumulate(weight, up, down, scale, "cpu")
 
 
-def _delta(up: Any, down: Any, weight: Any, device: Any) -> Any:
-    """``up @ down`` on ``device``, shaped and typed for the target weight. Conv LoRAs flatten the
-    spatial dims."""
+def _accumulate(weight: Any, up: Any, down: Any, scale: float, device: Any) -> None:
+    """Add the product into ``weight`` in row slices. Conv LoRAs flatten the spatial dims."""
     dtype = _fuse_dtype(up)
-    delta = up.to(device, dtype=dtype).flatten(1) @ down.to(device, dtype=dtype).flatten(1)
-    return delta.reshape(weight.shape).to(weight.dtype)
+    rows = up.to(device, dtype=dtype).flatten(1)
+    cols = down.to(device, dtype=dtype).flatten(1)
+    if not weight.is_contiguous():  # a non-contiguous target cannot be written through a view
+        weight.add_((rows @ cols).reshape(weight.shape).to(weight.dtype) * scale)
+        return
+    target = weight.view(rows.shape[0], -1)
+    step = max(1, _DELTA_CHUNK_BYTES // max(1, cols.shape[1] * 4))
+    for start in range(0, rows.shape[0], step):
+        stop = start + step
+        target[start:stop].add_((rows[start:stop] @ cols).to(weight.dtype) * scale)
 
 
 def _fuse_dtype(tensor: Any) -> Any:
@@ -141,20 +174,29 @@ def _alpha_scale(alpha: Any, rank: int) -> float:
     return float(alpha.item() if hasattr(alpha, "item") else alpha) / float(rank)
 
 
+def split_key(key: str) -> tuple[str, str] | None:
+    """``…to_q.lora_A.weight`` to ``("…to_q", "down")``. None for anything that is not a LoRA key.
+
+    Public so a per-arch key translator groups by exactly the suffixes the fuser recognises: a
+    convention known to one and not the other would drop tensors silently."""
+    for suffix in _DOWN:
+        if key.endswith("." + suffix):
+            return key[: -len(suffix) - 1], "down"
+    for suffix in _UP:
+        if key.endswith("." + suffix):
+            return key[: -len(suffix) - 1], "up"
+    if key.endswith(".alpha"):
+        return key[: -len(".alpha")], "alpha"
+    return None
+
+
 def _group(state: dict[str, Any]) -> tuple[dict[str, tuple[Any, Any]], dict[str, Any]]:
-    downs: dict[str, Any] = {}
-    ups: dict[str, Any] = {}
-    alphas: dict[str, Any] = {}
+    parts: dict[str, dict[str, Any]] = {}
     for key, value in state.items():
-        for suffix in _DOWN:
-            if key.endswith("." + suffix):
-                downs[key[: -len(suffix) - 1]] = value
-        for suffix in _UP:
-            if key.endswith("." + suffix):
-                ups[key[: -len(suffix) - 1]] = value
-        if key.endswith(".alpha"):
-            alphas[key[: -len(".alpha")]] = value
-    return {k: (downs[k], ups[k]) for k in downs if k in ups}, alphas
+        if (split := split_key(key)) is not None:
+            parts.setdefault(split[0], {})[split[1]] = value
+    pairs = {k: (v["down"], v["up"]) for k, v in parts.items() if "down" in v and "up" in v}
+    return pairs, {k: v["alpha"] for k, v in parts.items() if "alpha" in v}
 
 
 def _linear_module_names(model: Any) -> dict[str, None]:

{inline_core-1.2.66 → inline_core-1.2.68}/src/inline_core/models/minimaxh3/adaln.py

@@ -41,6 +41,11 @@ from typing import Any
 import torch
 from torch import nn
 
+from .vendor.transformer_minimax_h3 import (
+    MiniMaxH3AdaLayerNormModulation,
+    MiniMaxH3AdaLayerNormOut,
+)
+
 logger = logging.getLogger("inline_core.minimaxh3")
 
 #: Rank kept. Five directions carry the energy; eight matches the published build and leaves slack.
@@ -144,3 +149,100 @@ def factorise(model: Any, *, rank: int = RANK) -> int:
         len(model.transformer_blocks), saved / 1e9,
     )
     return saved
+
+
+# --- the published pruned builds ------------------------------------------------------------
+#
+# MiniMax ship `pruned` checkpoints that do this same rank-8 reduction ahead of time, and go one
+# step further: the timestep path itself is gone. There is no `time_embedder` in the file at all,
+# only `adaln_t_table [1025, 8]`, holding `silu(temb)` already projected into their basis at 1025
+# points across t in [0, 1]. So the branch cannot be rebuilt as a basis applied to a `silu(temb)` we
+# compute; the table has to be read directly.
+#
+# Off-grid timesteps are interpolated. Measured against the full bf16 weights at 24 random t, the
+# table reaches 1.636e-4 relative with linear interpolation and 1.874e-4 taking the nearest row,
+# where one bf16 ulp of the reference is 1.307e-3. The grid is dense enough that interpolating costs
+# nothing and removes the sampler constraint a lookup would otherwise impose.
+
+#: Rows in the published table. Checked, not assumed: a build on a different grid must not be read
+#: as though it were on this one.
+TABLE_ROWS = 1025
+
+
+class TableEmbedder(nn.Module):
+    """Stands in for ``time_proj`` + ``time_embedder``, returning the pruned build's rank-8 row.
+
+    ``linear_1`` exists because the port reads ``time_embedder.linear_1.weight.dtype`` to cast its
+    input; it carries the table's dtype and nothing else, which keeps ``vendor/`` verbatim.
+    """
+
+    #: Declared so the buffer reads as a tensor; ``register_buffer`` alone types as ``Module``.
+    table: torch.Tensor
+
+    def __init__(self, table: torch.Tensor) -> None:
+        super().__init__()
+        self.register_buffer("table", table, persistent=True)
+        self.linear_1 = nn.Linear(1, 1, bias=False, dtype=table.dtype)
+
+    def forward(self, timestep: torch.Tensor) -> torch.Tensor:
+        rows = self.table.shape[0]
+        position = timestep.to(self.table.dtype).flatten() * (rows - 1)
+        low = position.floor().long().clamp(0, rows - 2)
+        frac = (position - low).unsqueeze(1)
+        return self.table[low] * (1 - frac) + self.table[low + 1] * frac
+
+
+class TabulatedModulation(MiniMaxH3AdaLayerNormModulation):
+    """``adaln_proj`` reading a table row. It already holds ``silu(temb)``, so no activation."""
+
+    def forward(self, temb: torch.Tensor) -> tuple[torch.Tensor, ...]:
+        out = self.linear(temb.to(self.linear.weight.dtype)).view(-1, 6 * self.hidden_size)
+        return out.chunk(6, dim=-1)
+
+
+class TabulatedNormOut(MiniMaxH3AdaLayerNormOut):
+    """``norm_out`` reading a table row, otherwise the port's own forward."""
+
+    def forward(
+        self, hidden_states: torch.Tensor, temb: torch.Tensor, timestep_indices: torch.Tensor
+    ) -> torch.Tensor:
+        shift, scale = self.linear(temb.to(self.linear.weight.dtype)).chunk(2, dim=-1)
+        hidden_states = self.norm(hidden_states)
+        return hidden_states * (1.0 + scale.index_select(0, timestep_indices)) + shift.index_select(
+            0, timestep_indices
+        )
+
+
+@torch.no_grad()
+def tabulate(model: Any, table: torch.Tensor) -> None:
+    """Rebuild the timestep path around a pruned build's table, on the meta device before streaming.
+
+    Every module the table feeds changes shape, so this has to happen before any weight is placed:
+    the rank-8 ``adaln_proj`` in the file would otherwise be assigned into a ``[96768, 2688]`` slot.
+    """
+    if table.ndim != 2 or table.shape[0] != TABLE_ROWS:
+        raise ValueError(
+            f"adaln_t_table is {tuple(table.shape)}, not [{TABLE_ROWS}, rank]. This build is on a "
+            "different timestep grid from the one measured, and reading it as if it were not would "
+            "shift the modulation at every step while still rendering."
+        )
+    rank = int(table.shape[1])
+    model.time_proj = nn.Identity()
+    model.time_embedder = TableEmbedder(table)
+    for block in model.transformer_blocks:
+        block.adaln_proj = _retyped(block.adaln_proj, TabulatedModulation, rank)
+    model.norm_out = _retyped(model.norm_out, TabulatedNormOut, rank)
+
+
+def _retyped(module: Any, cls: type, rank: int) -> Any:
+    """The same module with its projection narrowed to ``rank`` inputs, still on meta."""
+    replacement = module
+    replacement.__class__ = cls
+    old = module.linear
+    with torch.device("meta"):
+        replacement.linear = nn.Linear(rank, old.out_features, bias=old.bias is not None)
+    return replacement
+
+
+#: Parameters ``tabulate`` creates that no checkpoint fills.
+TABULATED_SELF_COMPUTED = ("time_embedder.linear_1.weight", "time_embedder.table")

{inline_core-1.2.66 → inline_core-1.2.68}/src/inline_core/models/minimaxh3/keys.py

@@ -23,7 +23,7 @@ fused QKV into three.
 
 from __future__ import annotations
 
-from ..keymap import AssertEqual, KeyPlan, Rename, RowLayout, Split, SwapHalves
+from ..keymap import AssertEqual, Drop, KeyPlan, Rename, RowLayout, Split, SwapHalves
 
 #: Bumping this invalidates every prepared artifact built by the old plan.
 PLAN_VERSION = "minimax-h3.keys.1"
@@ -81,9 +81,17 @@ def build_plan(
     num_blocks: int = NUM_BLOCKS,
     num_refiner_blocks: int = NUM_REFINER_BLOCKS,
     head_dim: int = HEAD_DIM,
+    pruned: bool = False,
+    sidecars: tuple[str, ...] = (),
 ) -> KeyPlan:
     """The plan for a publisher's layout. ``source`` selects how the fused QKV rows are arranged.
 
+    ``pruned`` is the published rank-8 build: it ships ``adaln_t_table`` in place of the whole
+    timestep path, so the two ``time_embedder`` projections are simply not in the file.
+
+    ``sidecars`` are the quantisation tensors an fp8 build carries beside each weight. They are
+    consumed while streaming and dropped here, so the coverage check still accounts for every key.
+
     The counts are arguments so a round-trip test can exercise the same code at a size that fits in
     memory; the defaults are the released geometry.
     """
@@ -95,7 +103,16 @@ def build_plan(
         ) from None
 
     actions: dict[str, object] = {}
-    for stem, target in _TOP_LEVEL.items():
+    top_level = {
+        stem: target
+        for stem, target in _TOP_LEVEL.items()
+        if not (pruned and stem.startswith("time_embedder."))
+    }
+    if pruned:
+        actions["adaln_t_table"] = Drop("read before streaming, to rebuild the timestep path")
+    for key in sidecars:
+        actions[key] = Drop("a quantisation scale, applied to its weight while streaming")
+    for stem, target in top_level.items():
         for suffix in ("weight", "bias"):
             actions[f"{stem}.{suffix}"] = Rename(f"{target}.{suffix}")
     for key, target in _WEIGHT_ONLY.items():
@@ -126,9 +143,16 @@ def build_plan(
                     continue
                 actions[f"{src}.{stem}"] = Rename(f"{dst}.{target}")
 
-    return KeyPlan(version=f"{PLAN_VERSION}+{source}", actions=actions)  # type: ignore[arg-type]
+    suffix = source + ("+pruned" if pruned else "") + ("+fp8" if sidecars else "")
+    return KeyPlan(version=f"{PLAN_VERSION}+{suffix}", actions=actions)  # type: ignore[arg-type]
 
 
-def self_computed_targets() -> set[str]:
+def self_computed_targets(*, pruned: bool = False) -> set[str]:
     """Targets the port builds itself, which ``check_coverage`` must not demand be filled."""
-    return {"rope.inv_freq"}
+    if not pruned:
+        return {"rope.inv_freq"}
+    return {"rope.inv_freq", *_PRUNED_SELF_COMPUTED}
+
+
+#: What ``adaln.tabulate`` creates in place of the timestep path. No checkpoint fills these.
+_PRUNED_SELF_COMPUTED = ("time_embedder.table", "time_embedder.linear_1.weight")