inline-core 1.2.2__tar.gz → 1.2.31__tar.gz

{inline_core-1.2.2 → inline_core-1.2.31}/CLAUDE.md

@@ -82,8 +82,6 @@ HTTP/WS  →  server/app.py  →  RunManager  →  Executor  →  Registry (desc
 - **`runtime/`** - `context.py` (`ExecutionContext`, `CancelToken`), `run.py` (`RunState`),
   `progress.py` (events + emitters), `store.py` / `file_store.py` (`TakeStore`: owns take bytes/hash/
   uri).
-- **`importer/`** - `comfy.py`: best-effort ComfyUI-workflow → Inline-Core-graph mapping. All ComfyUI
-  format knowledge lives here.
 - **`config.py`** - all env config, small and explicit. **`takes.py`**, **`media.py`**, **`errors.py`**
   - domain primitives (`Take`/`AssetRef`, `MediaKind`, the error hierarchy).
 
@@ -200,13 +198,13 @@ real codec that moves tensors lives with the model runner.
   subpackages import torch/diffusers at module top **on purpose**: an absent extra makes the import
   raise, and `server/bootstrap.py` skips that model best-effort so a core install still boots and
   serves source nodes. Never import a heavy dep at package top level outside a runner subpackage.
-- **Engine isolation.** All ComfyUI-format knowledge lives in `importer/comfy.py`; all xDiT/worker
-  knowledge behind `parallel/` and the sampler seam. Don't scatter it.
+- **Engine isolation.** All xDiT/worker knowledge lives behind `parallel/` and the sampler seam.
+  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>"`.
 - **Tests (pytest).** Cover the logic that matters: graph validate/topo/executor/cache, the catalog
   scan, the run store + server contract, the device/memory policy, the parallel group + xfuser seam,
-  the comfy importer, and each model runner (import-guarded, no GPU needed). See `tests/`.
+  and each model runner (import-guarded, no GPU needed). See `tests/`.
 - **Commits.** Conventional Commits (`feat:`, `fix:`, `chore:`), small and scoped.
 
 ## Commands
@@ -239,4 +237,3 @@ uv run pytest -q                          # tests (no GPU; model code is import-
   as `{error:{code,message}}` with the right status. Update the API list in `README.md`.
 - **New port/handle type** → `PortKind` in `graph/schema.py` (+ `port_satisfies` if it has coercions).
 - **New device/memory behaviour** → behind `DevicePolicy` in `device/`; never in a component.
-- **New ComfyUI import behaviour** → `importer/comfy.py` only.

{inline_core-1.2.2 → inline_core-1.2.31}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: inline-core
-Version: 1.2.2
+Version: 1.2.31
 Summary: The generation engine behind Inline Studio.
 Requires-Python: >=3.11
 Requires-Dist: numpy>=1.26
@@ -53,7 +53,7 @@ already supports, so the multi-GPU split works on it from the start).
 > Status: early, and running end to end against a stub engine. In place and tested: the graph engine,
 > the typed `/v1` HTTP + websocket API (durable runs, streamed progress, coalescing), the model-dir
 > scan, the device + memory policy (profiles, dtype, offload, int8), the low-level primitive node
-> vocabulary, and a ComfyUI workflow importer. The Z-Image loader is written and validates on a GPU.
+> vocabulary. The Z-Image loader is written and validates on a GPU.
 > Cross-request batching, single-image multi-GPU (an xDiT worker group behind the sampler seam, with
 > the policy and IPC round-trip tested), and out-of-process custom nodes are built as seams but not
 > yet running on real hardware.
@@ -131,8 +131,7 @@ Underneath, a **low-level primitive vocabulary** exists (`load/diffusion-model`,
 `load/text-encoder`, `encode/text`, `latent/empty`, `sample`, `vae/decode`, `vae/encode`) - the
 ComfyUI-equivalent decomposed graph, kept for validation/execution - but these are marked **`hidden`**
 and never appear in the add-node menu. Engine handles (`model`, `vae`, `text-encoder`, `conditioning`,
-`latent`) are typed sockets between nodes; only media outputs become Frames with take history. A
-best-effort ComfyUI importer maps existing workflows onto the primitives.
+`latent`) are typed sockets between nodes; only media outputs become Frames with take history.
 
 ## Multi-GPU: split one image across GPUs
 

{inline_core-1.2.2 → inline_core-1.2.31}/README.md

@@ -11,7 +11,7 @@ already supports, so the multi-GPU split works on it from the start).
 > Status: early, and running end to end against a stub engine. In place and tested: the graph engine,
 > the typed `/v1` HTTP + websocket API (durable runs, streamed progress, coalescing), the model-dir
 > scan, the device + memory policy (profiles, dtype, offload, int8), the low-level primitive node
-> vocabulary, and a ComfyUI workflow importer. The Z-Image loader is written and validates on a GPU.
+> vocabulary. The Z-Image loader is written and validates on a GPU.
 > Cross-request batching, single-image multi-GPU (an xDiT worker group behind the sampler seam, with
 > the policy and IPC round-trip tested), and out-of-process custom nodes are built as seams but not
 > yet running on real hardware.
@@ -89,8 +89,7 @@ Underneath, a **low-level primitive vocabulary** exists (`load/diffusion-model`,
 `load/text-encoder`, `encode/text`, `latent/empty`, `sample`, `vae/decode`, `vae/encode`) - the
 ComfyUI-equivalent decomposed graph, kept for validation/execution - but these are marked **`hidden`**
 and never appear in the add-node menu. Engine handles (`model`, `vae`, `text-encoder`, `conditioning`,
-`latent`) are typed sockets between nodes; only media outputs become Frames with take history. A
-best-effort ComfyUI importer maps existing workflows onto the primitives.
+`latent`) are typed sockets between nodes; only media outputs become Frames with take history.
 
 ## Multi-GPU: split one image across GPUs
 

{inline_core-1.2.2 → inline_core-1.2.31}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 # PyPI name; the import package is `inline_core` (src/inline_core).
 name = "inline-core"
-version = "1.2.2"
+version = "1.2.31"
 description = "The generation engine behind Inline Studio."
 readme = "README.md"
 requires-python = ">=3.11"

{inline_core-1.2.2 → inline_core-1.2.31}/src/inline_core/extensions/fetch.py

@@ -49,12 +49,18 @@ def remote_sha(url: str, ref: str) -> str | None:
     an update check must never fail the dialog."""
     try:
         _validate(url, ref)
-        done = _git("ls-remote", url, ref, check=False)
+        # Also request the peeled ref: an exact-match pattern alone won't return the `^{}` line.
+        done = _git("ls-remote", url, ref, f"{ref}^{{}}", check=False)
     except (FetchError, OSError, subprocess.SubprocessError):
         return None
     if done.returncode != 0:
         return None
-    line = done.stdout.strip().split("\n")[0]
+    lines = [ln for ln in done.stdout.splitlines() if ln.strip()]
+    # An annotated tag lists its own object sha first and the commit it points to on a `^{}` line.
+    # Prefer the peeled commit, since the installed sha is a commit (`_resolve` uses `^{commit}`) -
+    # otherwise every annotated-tag release reads as perpetually "update available".
+    peeled = next((ln for ln in lines if ln.rstrip().endswith("^{}")), None)
+    line = peeled or (lines[0] if lines else "")
     sha = line.split()[0] if line else ""
     return sha or None
 

inline_core-1.2.31/src/inline_core/ffmpeg.py

@@ -0,0 +1,30 @@
+"""Locate the ffmpeg/ffprobe binaries. Prefers a bundled ``imageio-ffmpeg``, else PATH.
+
+Lives at the top level rather than under ``studio/`` because both the timeline (studio) and the
+take store (runtime) need it, and runtime must not import studio.
+"""
+
+from __future__ import annotations
+
+import shutil
+from functools import lru_cache
+
+
+@lru_cache(maxsize=1)
+def ffmpeg_exe() -> str | None:
+    try:
+        import imageio_ffmpeg
+
+        return imageio_ffmpeg.get_ffmpeg_exe()
+    except Exception:  # noqa: BLE001
+        return shutil.which("ffmpeg")
+
+
+@lru_cache(maxsize=1)
+def ffprobe_exe() -> str | None:
+    """PATH only - imageio bundles ffmpeg alone, so probing degrades gracefully when absent."""
+    return shutil.which("ffprobe")
+
+
+def ffmpeg_available() -> bool:
+    return ffmpeg_exe() is not None

{inline_core-1.2.2 → inline_core-1.2.31}/src/inline_core/graph/loader_runners.py

@@ -20,7 +20,7 @@ from typing import TYPE_CHECKING, Any
 
 from ..config import models_dir
 from ..errors import ComponentError
-from .primitives import LOAD_DIFFUSION_MODEL, LOAD_TEXT_ENCODER, LOAD_VAE
+from .primitives import LOAD_DIFFUSION_MODEL, LOAD_LORA, LOAD_TEXT_ENCODER, LOAD_VAE
 from .runners import NodeResult, NodeRunner
 from .schema import Node
 
@@ -44,6 +44,16 @@ class ComponentRef:
     file: str
 
 
+@dataclass(frozen=True)
+class LoraRef:
+    """One LoRA in a stack: its absolute file path and blend strength. A ``load/lora`` node emits a
+    tuple of these (its own ref appended to any upstream stack); the model runner fuses them into
+    the diffusion transformer in order. Frozen + hashable so it can key the loader cache."""
+
+    file: str
+    strength: float
+
+
 def _resolve_file(category: str, chosen: str) -> Path:
     """The single weight file a Load node points at: the explicit dropdown pick, else the first
     weight file in ``models/<category>/`` (mirrors the model node's "auto"). Raises if none."""
@@ -81,6 +91,26 @@ class LoadComponentRunner(NodeRunner):
         return NodeResult(outputs={self._output: ref})
 
 
+class LoadLoraRunner(NodeRunner):
+    """Resolve this node's ``file``/``strength`` into a ``LoraRef`` and append it to any upstream
+    stack on the ``lora`` input - so chaining ``load/lora`` nodes stacks them in wiring order. The
+    stack rides its own ``lora`` edge into the model runner, which fuses it into the transformer."""
+
+    produces_takes = False
+
+    def run(self, node: Node, inputs: dict[str, list[Any]], ctx: ExecutionContext) -> NodeResult:
+        upstream = _first(inputs.get("lora")) or ()
+        file = _resolve_file("loras", str(node.params.get("file", "")))
+        strength = float(node.params.get("strength", 1.0))
+        stack = (*upstream, LoraRef(file=str(file), strength=strength))
+        return NodeResult(outputs={"lora": stack})
+
+
+def _first(values: list[Any] | None) -> Any:
+    """The first wired value on a port, or None (an optional input may be absent/unconnected)."""
+    return values[0] if values else None
+
+
 def register_loaders(registry: Registry) -> None:
     """Register the ``load/*`` nodes **visible** (unhidden) with their runners, so they appear in
     the add-node menu and can feed a model node's component inputs. Torch-free - always on."""
@@ -98,3 +128,4 @@ def register_loaders(registry: Registry) -> None:
             kind="text_encoder", category="text_encoders", output_port="text_encoder"
         ),
     )
+    registry.register(replace(LOAD_LORA, hidden=False), LoadLoraRunner())

{inline_core-1.2.2 → inline_core-1.2.31}/src/inline_core/graph/primitives.py

@@ -47,6 +47,20 @@ LOAD_TEXT_ENCODER = NodeDescriptor(
     outputs=(Port("text_encoder", "Text encoder", PortKind.TEXT_ENCODER),),
 )
 
+LOAD_LORA = NodeDescriptor(
+    type="load/lora",
+    title="Load LoRA",
+    category="Loaders",
+    icon="box",
+    # Optional upstream lora: chain load/lora -> load/lora to stack several in order.
+    inputs=(Port("lora", "LoRA", PortKind.LORA, required=False),),
+    params=(
+        ParamField("file", "LoRA", Widget.SELECT, "", options_from="loras"),
+        ParamField("strength", "Strength", Widget.NUMBER, 1.0, min=-2.0, max=2.0, step=0.05),
+    ),
+    outputs=(Port("lora", "LoRA", PortKind.LORA),),
+)
+
 ENCODE_TEXT = NodeDescriptor(
     type="encode/text",
     title="Encode Text",
@@ -122,6 +136,7 @@ PRIMITIVES: tuple[NodeDescriptor, ...] = (
     LOAD_DIFFUSION_MODEL,
     LOAD_VAE,
     LOAD_TEXT_ENCODER,
+    LOAD_LORA,
     ENCODE_TEXT,
     EMPTY_LATENT,
     SAMPLE,
@@ -132,7 +147,7 @@ PRIMITIVES: tuple[NodeDescriptor, ...] = (
 
 #: The loader primitives now have runners and are offered in the add-node menu - registered
 #: (unhidden) with their runners by ``graph/loader_runners.py``, so they are skipped here.
-_HAS_RUNNER = {LOAD_DIFFUSION_MODEL.type, LOAD_VAE.type, LOAD_TEXT_ENCODER.type}
+_HAS_RUNNER = {LOAD_DIFFUSION_MODEL.type, LOAD_VAE.type, LOAD_TEXT_ENCODER.type, LOAD_LORA.type}
 
 
 def register_primitives(registry: Registry) -> None:

{inline_core-1.2.2 → inline_core-1.2.31}/src/inline_core/graph/schema.py

@@ -23,6 +23,7 @@ class PortKind(str, Enum):
     MODEL = "model"
     VAE = "vae"
     TEXT_ENCODER = "text-encoder"
+    LORA = "lora"
     CONDITIONING = "conditioning"
     LATENT = "latent"
 

{inline_core-1.2.2 → inline_core-1.2.31}/src/inline_core/models/loaders.py

@@ -30,12 +30,15 @@ from collections.abc import Callable
 from dataclasses import dataclass
 from pathlib import Path
 from threading import Lock
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from ..config import data_dir
 from ..device.policy import Quantization
 from ..errors import ComponentError
 
+if TYPE_CHECKING:
+    from ..graph.loader_runners import LoraRef
+
 
 @dataclass(frozen=True)
 class ArchSpec:
@@ -260,13 +263,15 @@ def _quantize_in_place(model: Any, quant: Quantization) -> None:
 # --- component loaders (cached in-process) ------------------------------------------------------
 
 # Keyed by (arch, kind, file, dtype, quant, device) so switching one file (e.g. a different VAE),
-# the quantization, or the load device reuses the other already-loaded components. The run manager
-# executes one run at a time; the lock guards each build.
-_CACHE: dict[tuple[str, str, str, str, str, str], Any] = {}
+# the quantization, or the load device reuses the other already-loaded components. A fused diffusion
+# transformer appends its LoRA stack (see ``lora_cache_key``) to that base key, so the same file
+# with a different stack is a distinct entry. The run manager executes one run at a time; the lock
+# guards each build.
+_CACHE: dict[tuple[str, ...], Any] = {}
 _CACHE_LOCK = Lock()
 
 
-def _cached(key: tuple[str, str, str, str, str, str], build: Callable[[], Any]) -> Any:
+def _cached(key: tuple[str, ...], build: Callable[[], Any]) -> Any:
     with _CACHE_LOCK:
         hit = _CACHE.get(key)
         if hit is not None:
@@ -284,17 +289,36 @@ def _device_key(device: str | None) -> str:
     return device or "cpu"
 
 
-def unload_components(keep_files: set[str] | None = None) -> None:
+def lora_cache_key(loras: tuple[LoraRef, ...]) -> tuple[str, ...]:
+    """The cache-key suffix a LoRA stack contributes to a fused diffusion transformer. Order- and
+    strength-sensitive (fusing is not commutative), keyed by file+strength; an empty stack adds
+    nothing, so an un-LoRA'd load keys exactly as it did before LoRA existed."""
+    return tuple(f"{lora.file}@{lora.strength}" for lora in loras)
+
+
+def unload_components(
+    keep_files: set[str] | None = None, keep_loras: tuple[str, ...] | None = None
+) -> None:
     """Drop cached components whose source file is NOT in ``keep_files``, freeing their VRAM/RAM.
 
     Called when switching checkpoints so a new model doesn't stack on top of the previous one (the
     cache never evicted before, so a second distinct model roughly doubled resident memory). Only
     drops references + empties the CUDA cache - it does not move weights to CPU RAM (that would just
     relocate the pressure on a RAM-tight box). The caller must drop any pipeline holding these
-    components first, or the references keep them alive."""
+    components first, or the references keep them alive.
+
+    ``keep_loras`` (the LoRA-stack suffix being loaded, from ``lora_cache_key``) additionally evicts
+    a kept file's diffusion transformer when it carries a *different* stack: fusing a LoRA into an
+    already-resident checkpoint would otherwise keep the unfused transformer AND the fused one - two
+    full-size models on the card. Left None, eviction is file-only (the pre-LoRA behaviour)."""
     keep = keep_files or set()
     with _CACHE_LOCK:
-        stale = [k for k in _CACHE if k[2] not in keep]  # k = (arch, kind, file, dtype, quant, dev)
+        stale = []
+        for k in _CACHE:  # k = (arch, kind, file, dtype, quant, dev, *lora_suffix)
+            if k[2] not in keep:
+                stale.append(k)
+            elif keep_loras is not None and k[1] == "diffusion" and tuple(k[6:]) != keep_loras:
+                stale.append(k)
         for k in stale:
             comp = _CACHE.pop(k)
             del comp
@@ -307,6 +331,7 @@ def load_diffusion(
     dtype: Any,
     quant: Quantization = Quantization.NONE,
     device: str | None = None,
+    loras: tuple[LoraRef, ...] = (),
 ) -> Any:
     """The diffusion transformer from a single ``.safetensors``. diffusers converts the checkpoint
     keys; the config comes from the bundled assets, so nothing is fetched at load time. ``quant``
@@ -314,7 +339,8 @@ def load_diffusion(
     **straight to the GPU** from an mmap-backed checkpoint - the fp16 weights are never materialized
     as an anonymous CPU copy (the host-RAM spike that OOM-killed the server), and for the int8 path
     torchao quantizes on-device per tensor. ``None`` loads to CPU (the offload path, where
-    accelerate installs its hooks before placing)."""
+    accelerate installs its hooks before placing). ``loras`` are fused into the weights in order
+    (the stack is part of the cache key), so a fused transformer is itself the cached artifact."""
 
     def build() -> Any:
         from diffusers import ZImageTransformer2DModel
@@ -333,12 +359,20 @@ def load_diffusion(
         # ``from_pretrained``), so the transformer would load at full size and the "int8" plan would
         # blow the VRAM budget (a T4 OOMs mid-load). Quantize it explicitly with torchao after the
         # load instead - the weights briefly sit full-size on the device, then halve in place.
+        # Fuse LoRAs BEFORE quantizing: the fuse adds a full-precision delta into each weight, which
+        # int8 can't accept in place, and int8-quantized weights aren't a plain tensor to add into.
+        if loras:
+            from .lora import fuse_loras
+
+            fuse_loras(model, loras)
         _quantize_in_place(model, quant)
         return model
 
-    return _cached(
-        (arch, "diffusion", file, _dtype_key(dtype), quant.value, _device_key(device)), build
+    key = (
+        arch, "diffusion", file, _dtype_key(dtype), quant.value, _device_key(device),
+        *lora_cache_key(loras),
     )
+    return _cached(key, build)
 
 
 def load_vae(arch: str, file: str, dtype: Any, device: str | None = None) -> Any:
@@ -431,6 +465,7 @@ def assemble_zimage_pipeline(
     quant: Quantization = Quantization.NONE,
     vae_dtype: Any = None,
     device: str | None = None,
+    loras: tuple[LoraRef, ...] = (),
     cancel_check: Callable[[], None] | None = None,
 ) -> Any:
     """Build a Z-Image pipeline from three local single files. Components are cached individually,
@@ -448,7 +483,7 @@ def assemble_zimage_pipeline(
     from diffusers import ZImageImg2ImgPipeline, ZImagePipeline
 
     arch = _ZIMAGE.key
-    transformer = load_diffusion(arch, diffusion_file, dtype, quant, device=device)
+    transformer = load_diffusion(arch, diffusion_file, dtype, quant, device=device, loras=loras)
     _release_transient()
     if cancel_check is not None:
         cancel_check()

inline_core-1.2.31/src/inline_core/models/lora.py

@@ -0,0 +1,140 @@
+"""Fuse LoRA weights into an already-loaded model, in stack order.
+
+Fusing (rather than keeping adapters live) is deliberate: ``loaders.py`` caches components by a key
+that includes the stack, so a fused model *is* the cached artifact and there is no adapter state to
+get out of sync with that key.
+
+Handles both key conventions we see in the wild - ComfyUI/kohya (``lora_down``/``lora_up`` + an
+optional ``alpha``) and diffusers/peft (``lora_A``/``lora_B``). Every LoRA key must resolve to a
+module: a partial match silently applies some of the LoRA and no-ops the rest, which degrades output
+without erroring, so an unmatched key is a hard failure.
+
+Known gap: fully underscore-flattened kohya paths (``lora_unet_layers_0_attention_to_q``) are
+ambiguous - the reverse mapping cannot tell a path separator from an underscore inside a name like
+``to_q``. Those raise rather than mis-apply. Every Z-Image LoRA seen so far uses dotted paths.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from ..errors import ComponentError
+
+if TYPE_CHECKING:
+    from ..graph.loader_runners import LoraRef
+
+_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.
+_PREFIXES = ("diffusion_model.", "transformer.", "lora_unet_", "lora_te_", "base_model.model.")
+
+
+def fuse_loras(model: Any, loras: tuple[LoraRef, ...]) -> None:
+    """Merge each LoRA into ``model``'s weights in order. No-op for an empty stack."""
+    for lora in loras:
+        _fuse_one(model, lora.file, lora.strength)
+
+
+def _fuse_one(model: Any, path: str, strength: float) -> None:
+    import torch
+    from safetensors.torch import load_file
+
+    try:
+        state = load_file(path)
+    except Exception as exc:  # noqa: BLE001
+        raise ComponentError(f"Could not read LoRA {path!r}: {exc}") from exc
+
+    modules = _linear_modules(model)
+    pairs, alphas = _group(state)
+    if not pairs:
+        raise ComponentError(f"LoRA {path!r} contains no recognisable lora_down/lora_up pairs.")
+
+    unmatched: list[str] = []
+    fused = 0
+    for stem, (down, up) in sorted(pairs.items()):
+        target = _match(stem, modules)
+        if target is None:
+            unmatched.append(stem)
+            continue
+        with torch.no_grad():
+            weight = target.weight
+            delta = _delta(up, down, weight)
+            scale = strength * _alpha_scale(alphas.get(stem), down.shape[0])
+            weight.add_((delta * scale).to(device=weight.device, dtype=weight.dtype))
+        fused += 1
+
+    if unmatched:
+        sample = ", ".join(unmatched[:3])
+        raise ComponentError(
+            f"LoRA {path!r} does not match this model: {len(unmatched)} of {len(pairs)} layers "
+            f"have no target (e.g. {sample}). It was probably trained for a different architecture."
+        )
+    if fused == 0:
+        raise ComponentError(f"LoRA {path!r} matched no layers in this model.")
+
+
+def _delta(up: Any, down: Any, weight: Any) -> Any:
+    """``up @ down``, shaped to the target weight. Conv LoRAs flatten the spatial dims."""
+    dtype = _fuse_dtype(up)
+    delta = up.to("cpu", dtype=dtype).flatten(1) @ down.to("cpu", dtype=dtype).flatten(1)
+    return delta.reshape(weight.shape)
+
+
+def _fuse_dtype(tensor: Any) -> Any:
+    import torch
+
+    # fp16 matmuls of low-rank factors lose precision; fuse in fp32 and cast on the way in.
+    return torch.float32 if tensor.dtype in (torch.float16, torch.bfloat16) else tensor.dtype
+
+
+def _alpha_scale(alpha: Any, rank: int) -> float:
+    """kohya-style checkpoints scale by ``alpha / rank``; without an alpha the factors are raw."""
+    if alpha is None or rank == 0:
+        return 1.0
+    return float(alpha.item() if hasattr(alpha, "item") else alpha) / float(rank)
+
+
+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] = {}
+    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
+
+
+def _linear_modules(model: Any) -> dict[str, Any]:
+    import torch
+
+    return {
+        name: module
+        for name, module in model.named_modules()
+        if isinstance(module, torch.nn.Linear | torch.nn.Conv2d | torch.nn.Conv3d)
+    }
+
+
+def _match(stem: str, modules: dict[str, Any]) -> Any:
+    """Resolve a checkpoint's module path to a module, tolerating the usual prefixes and the
+    underscore-flattened paths kohya emits."""
+    for candidate in _candidates(stem):
+        hit = modules.get(candidate)
+        if hit is not None:
+            return hit
+    return None
+
+
+def _candidates(stem: str) -> list[str]:
+    out = [stem]
+    for prefix in _PREFIXES:
+        if stem.startswith(prefix):
+            out.append(stem[len(prefix) :])
+    # kohya flattens the whole path with underscores; the dotted form is the model's own naming.
+    out.extend([c.replace("_", ".") for c in list(out) if "." not in c])
+    return out