add LoRA loader node and weight fusion

Author: imprsnst

README.md

@@ -7,7 +7,7 @@
 <p align="center">
   <a href="LICENSE"><img alt="License: MIT" src="https://img.shields.io/badge/License-MIT-green?style=for-the-badge"></a>
   <a href="https://www.python.org/downloads/"><img alt="Python 3.11+" src="https://img.shields.io/badge/Python-3.11%2B-blue?style=for-the-badge&logo=python&logoColor=white"></a>
-  <a href="../../releases/latest"><img alt="Latest release" src="https://img.shields.io/badge/Release-v1.2.3-blue?style=for-the-badge"></a>
+  <a href="../../releases/latest"><img alt="Latest release" src="https://img.shields.io/badge/Release-v1.2.31-blue?style=for-the-badge"></a>
   <a href="https://discord.gg/cSUS88VdY9"><img alt="Join our Discord" src="https://img.shields.io/badge/Discord-Join%20the%20community-5865F2?logo=discord&logoColor=white&style=for-the-badge"></a>
 </p>
 

core/pyproject.toml

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

core/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

core/src/inline_core/graph/loader_runners.py

@@ -20,7 +20,7 @@
 
 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 @@ def run(self, node: Node, inputs: dict[str, list[Any]], ctx: ExecutionContext) -
         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())

core/src/inline_core/graph/primitives.py

@@ -47,6 +47,20 @@
     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 @@
     LOAD_DIFFUSION_MODEL,
     LOAD_VAE,
     LOAD_TEXT_ENCODER,
+    LOAD_LORA,
     ENCODE_TEXT,
     EMPTY_LATENT,
     SAMPLE,
@@ -132,7 +147,7 @@
 
 #: 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:

core/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"
 

core/src/inline_core/models/loaders.py

@@ -30,12 +30,15 @@
 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,14 +331,16 @@ 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``
     (smart memory) quantizes the weights on load. ``device`` (e.g. ``"cuda:0"``) streams each tensor
     **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 build() -> Any:
         # ``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()