inline-core 1.2.3__tar.gz → 1.2.5__tar.gz

{inline_core-1.2.3 → inline_core-1.2.5}/PKG-INFO

@@ -1,13 +1,14 @@
 Metadata-Version: 2.4
 Name: inline-core
-Version: 1.2.3
+Version: 1.2.5
 Summary: The generation engine behind Inline Studio.
+License-Expression: GPL-3.0-or-later
 Requires-Python: >=3.11
 Requires-Dist: numpy>=1.26
 Requires-Dist: psutil>=5.9
 Provides-Extra: all
 Requires-Dist: accelerate>=0.30; extra == 'all'
-Requires-Dist: diffusers>=0.36; extra == 'all'
+Requires-Dist: diffusers>=0.39; extra == 'all'
 Requires-Dist: fastapi>=0.110; extra == 'all'
 Requires-Dist: huggingface-hub>=0.23; extra == 'all'
 Requires-Dist: imageio-ffmpeg>=0.4; extra == 'all'
@@ -27,7 +28,7 @@ Requires-Dist: nvidia-ml-py>=12; extra == 'parallel'
 Requires-Dist: xfuser>=0.4; extra == 'parallel'
 Provides-Extra: runtime
 Requires-Dist: accelerate>=0.30; extra == 'runtime'
-Requires-Dist: diffusers>=0.36; extra == 'runtime'
+Requires-Dist: diffusers>=0.39; extra == 'runtime'
 Requires-Dist: huggingface-hub>=0.23; extra == 'runtime'
 Requires-Dist: safetensors>=0.4; extra == 'runtime'
 Requires-Dist: scipy>=1.11; extra == 'runtime'
@@ -38,6 +39,14 @@ Provides-Extra: server
 Requires-Dist: fastapi>=0.110; extra == 'server'
 Requires-Dist: imageio-ffmpeg>=0.4; extra == 'server'
 Requires-Dist: uvicorn[standard]>=0.29; extra == 'server'
+Provides-Extra: training
+Requires-Dist: bitsandbytes>=0.43; (platform_system != 'Darwin') and extra == 'training'
+Requires-Dist: einops>=0.7; extra == 'training'
+Requires-Dist: nvidia-ml-py>=12; extra == 'training'
+Requires-Dist: peft>=0.11; extra == 'training'
+Requires-Dist: pillow>=10; extra == 'training'
+Requires-Dist: psutil>=5.9; extra == 'training'
+Requires-Dist: timm>=1.0; extra == 'training'
 Description-Content-Type: text/markdown
 
 # Inline Core
@@ -47,8 +56,9 @@ The generation engine behind Inline. It takes a typed node graph (JSON) and retu
 low-VRAM laptops up to multi-GPU machines that split a single image's sampling across GPUs (via
 xDiT). It is Inline Studio's built-in render backend.
 
-First model: Z-Image (Alibaba Tongyi), a 6B rectified-flow diffusion transformer (and a model xDiT
-already supports, so the multi-GPU split works on it from the start).
+Models today: **Z-Image** (Alibaba Tongyi), a 6B rectified-flow diffusion transformer and the model
+xDiT already supports, so the multi-GPU split works on it from the start; and **Krea 2** (Krea AI),
+a 12.9B single-stream MMDiT released as an undistilled RAW base plus an 8-step distilled Turbo.
 
 > 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
@@ -94,9 +104,13 @@ by category:
 
 ```
 models/
-  diffusion_models/   z_image_turbo_bf16.safetensors   <- the one file you need for Z-Image
-  vae/                ae.safetensors                   <- optional (see below)
-  text_encoders/      qwen/   (an HF-format folder: config + tokenizer + weights)  <- optional
+  diffusion_models/   z_image_turbo_bf16.safetensors        <- Z-Image
+                      krea2_turbo_bf16.safetensors          <- Krea 2 Turbo (generate)
+                      krea2_raw_bf16.safetensors            <- Krea 2 RAW (fine-tune)
+  vae/                ae.safetensors                        <- Z-Image
+                      qwen_image_vae_diffusers.safetensors  <- Krea 2
+  text_encoders/      qwen_3_4b.safetensors                 <- Z-Image
+                      qwen3vl_4b_bf16.safetensors           <- Krea 2
   loras/  controlnet/  checkpoints/  ...
 ```
 
@@ -115,8 +129,17 @@ paths:
    `models/`** (never the hidden HF cache) with visible progress.
 
 Override the diffusion source with `INLINE_ZIMAGE_MODEL` (a file or a diffusers dir), and the supporting
-components with `INLINE_ZIMAGE_VAE` / `INLINE_ZIMAGE_TEXT_ENCODER`. The engine scans the models dir on
-start; a node's model pickers list what is present.
+components with `INLINE_ZIMAGE_VAE` / `INLINE_ZIMAGE_TEXT_ENCODER`; Krea 2 has the same three under
+`INLINE_KREA2_*`. The engine scans the models dir on start; a node's model pickers list what is present.
+
+**Krea 2 loads from the same ComfyUI single files**, with two constraints worth knowing. Only the
+**bf16** builds in [`Comfy-Org/Krea-2`](https://huggingface.co/Comfy-Org/Krea-2) are loadable: the
+fp8/int8/nvfp4 variants carry ComfyUI-specific scale tensors, and the loader refuses them by name
+rather than failing mid-load. And the **VAE must be the diffusers-format** Qwen-Image file (the
+node's popup fetches it from `Qwen/Qwen-Image`), because ComfyUI's copy of the same weights uses a
+module layout diffusers has no converter for. `Krea2Transformer2DModel` has no `from_single_file`, so
+the transformer checkpoint is renamed into diffusers naming on load (`models/krea2/convert.py`) and
+streamed tensor by tensor onto the GPU.
 
 ## Nodes
 
@@ -124,8 +147,9 @@ start; a node's model pickers list what is present.
 renders any node generically - adding a node type needs no UI release.
 
 **High-level model nodes are what the user sees.** Generation is one-click: you drop a single
-**Z-Image Turbo** node, wire a Prompt into it, and hit Run. The node hooks up the diffusion model, VAE,
-and text-encoder behind the scenes - no loader/sampler wiring.
+**Z-Image Turbo**, **Krea 2 Turbo** or **Krea 2 RAW** node, wire a Prompt into it, and hit Run. The
+node hooks up the diffusion model, VAE, and text-encoder behind the scenes - no loader/sampler
+wiring.
 
 Underneath, a **low-level primitive vocabulary** exists (`load/diffusion-model`, `load/vae`,
 `load/text-encoder`, `encode/text`, `latent/empty`, `sample`, `vae/decode`, `vae/encode`) - the
@@ -214,7 +238,7 @@ The easy path is `webui.sh`, which maps friendly flags onto the engine's `INLINE
 ./webui.sh --listen --port 9000       # bind all interfaces on 9000
 ./webui.sh --multi-gpu                # split one image across GPUs (auto with 2+ GPUs)
 ./webui.sh --lowvram                  # tight-VRAM profile
-./webui.sh --install --extra zimage   # set up ./.venv with the Z-Image runtime, then exit
+./webui.sh --install --extra runtime  # set up ./.venv with the model runtime, then exit
 ```
 
 `./webui.sh --help` lists every flag (networking, multi-GPU, device/memory profile, paths). The same

{inline_core-1.2.3 → inline_core-1.2.5}/README.md

@@ -5,8 +5,9 @@ The generation engine behind Inline. It takes a typed node graph (JSON) and retu
 low-VRAM laptops up to multi-GPU machines that split a single image's sampling across GPUs (via
 xDiT). It is Inline Studio's built-in render backend.
 
-First model: Z-Image (Alibaba Tongyi), a 6B rectified-flow diffusion transformer (and a model xDiT
-already supports, so the multi-GPU split works on it from the start).
+Models today: **Z-Image** (Alibaba Tongyi), a 6B rectified-flow diffusion transformer and the model
+xDiT already supports, so the multi-GPU split works on it from the start; and **Krea 2** (Krea AI),
+a 12.9B single-stream MMDiT released as an undistilled RAW base plus an 8-step distilled Turbo.
 
 > 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
@@ -52,9 +53,13 @@ by category:
 
 ```
 models/
-  diffusion_models/   z_image_turbo_bf16.safetensors   <- the one file you need for Z-Image
-  vae/                ae.safetensors                   <- optional (see below)
-  text_encoders/      qwen/   (an HF-format folder: config + tokenizer + weights)  <- optional
+  diffusion_models/   z_image_turbo_bf16.safetensors        <- Z-Image
+                      krea2_turbo_bf16.safetensors          <- Krea 2 Turbo (generate)
+                      krea2_raw_bf16.safetensors            <- Krea 2 RAW (fine-tune)
+  vae/                ae.safetensors                        <- Z-Image
+                      qwen_image_vae_diffusers.safetensors  <- Krea 2
+  text_encoders/      qwen_3_4b.safetensors                 <- Z-Image
+                      qwen3vl_4b_bf16.safetensors           <- Krea 2
   loras/  controlnet/  checkpoints/  ...
 ```
 
@@ -73,8 +78,17 @@ paths:
    `models/`** (never the hidden HF cache) with visible progress.
 
 Override the diffusion source with `INLINE_ZIMAGE_MODEL` (a file or a diffusers dir), and the supporting
-components with `INLINE_ZIMAGE_VAE` / `INLINE_ZIMAGE_TEXT_ENCODER`. The engine scans the models dir on
-start; a node's model pickers list what is present.
+components with `INLINE_ZIMAGE_VAE` / `INLINE_ZIMAGE_TEXT_ENCODER`; Krea 2 has the same three under
+`INLINE_KREA2_*`. The engine scans the models dir on start; a node's model pickers list what is present.
+
+**Krea 2 loads from the same ComfyUI single files**, with two constraints worth knowing. Only the
+**bf16** builds in [`Comfy-Org/Krea-2`](https://huggingface.co/Comfy-Org/Krea-2) are loadable: the
+fp8/int8/nvfp4 variants carry ComfyUI-specific scale tensors, and the loader refuses them by name
+rather than failing mid-load. And the **VAE must be the diffusers-format** Qwen-Image file (the
+node's popup fetches it from `Qwen/Qwen-Image`), because ComfyUI's copy of the same weights uses a
+module layout diffusers has no converter for. `Krea2Transformer2DModel` has no `from_single_file`, so
+the transformer checkpoint is renamed into diffusers naming on load (`models/krea2/convert.py`) and
+streamed tensor by tensor onto the GPU.
 
 ## Nodes
 
@@ -82,8 +96,9 @@ start; a node's model pickers list what is present.
 renders any node generically - adding a node type needs no UI release.
 
 **High-level model nodes are what the user sees.** Generation is one-click: you drop a single
-**Z-Image Turbo** node, wire a Prompt into it, and hit Run. The node hooks up the diffusion model, VAE,
-and text-encoder behind the scenes - no loader/sampler wiring.
+**Z-Image Turbo**, **Krea 2 Turbo** or **Krea 2 RAW** node, wire a Prompt into it, and hit Run. The
+node hooks up the diffusion model, VAE, and text-encoder behind the scenes - no loader/sampler
+wiring.
 
 Underneath, a **low-level primitive vocabulary** exists (`load/diffusion-model`, `load/vae`,
 `load/text-encoder`, `encode/text`, `latent/empty`, `sample`, `vae/decode`, `vae/encode`) - the
@@ -172,7 +187,7 @@ The easy path is `webui.sh`, which maps friendly flags onto the engine's `INLINE
 ./webui.sh --listen --port 9000       # bind all interfaces on 9000
 ./webui.sh --multi-gpu                # split one image across GPUs (auto with 2+ GPUs)
 ./webui.sh --lowvram                  # tight-VRAM profile
-./webui.sh --install --extra zimage   # set up ./.venv with the Z-Image runtime, then exit
+./webui.sh --install --extra runtime  # set up ./.venv with the model runtime, then exit
 ```
 
 `./webui.sh --help` lists every flag (networking, multi-GPU, device/memory profile, paths). The same

{inline_core-1.2.3 → inline_core-1.2.5}/pyproject.toml

@@ -1,9 +1,10 @@
 [project]
 # PyPI name; the import package is `inline_core` (src/inline_core).
 name = "inline-core"
-version = "1.2.3"
+version = "1.2.5"
 description = "The generation engine behind Inline Studio."
 readme = "README.md"
+license = "GPL-3.0-or-later"
 requires-python = ">=3.11"
 # Deliberately torch-free so the engine boots and serves source nodes with no ML stack installed.
 dependencies = [
@@ -16,8 +17,8 @@ dependencies = [
 runtime = [
   # A bare `torch` is CPU-only on Windows; see [tool.uv] below and requirements.txt.
   "torch>=2.2",
-  # Z-Image needs diffusers >= 0.36 (ZImagePipeline).
-  "diffusers>=0.36",
+  # Krea 2 needs diffusers >= 0.39 (Krea2Pipeline); Z-Image needs >= 0.36 (ZImagePipeline).
+  "diffusers>=0.39",
   "transformers>=4.44",
   "accelerate>=0.30",
   "safetensors>=0.4",
@@ -39,6 +40,20 @@ parallel = [
   "xfuser>=0.4",
   "nvidia-ml-py>=12",
 ]
+# LoRA training (the Trainer tab): PEFT adapter training + local auto-caption + host/GPU telemetry.
+# Install alongside `runtime`: `.[runtime,training]`. Reuses runtime's torch/diffusers/accelerate.
+training = [
+  "peft>=0.11",
+  # 8-bit Adam keeps optimizer state small; no macOS wheels, so skip it there (AdamW fallback).
+  "bitsandbytes>=0.43; platform_system != 'Darwin'",
+  "Pillow>=10",
+  # Telemetry (CPU/RAM + per-GPU NVML) for the Trainer tab.
+  "psutil>=5.9",
+  "nvidia-ml-py>=12",
+  # Florence-2 auto-captioner is loaded via trust_remote_code and needs these.
+  "timm>=1.0",
+  "einops>=0.7",
+]
 dev = [
   "pytest>=8",
   "ruff>=0.6",
@@ -48,7 +63,7 @@ dev = [
 # Everything needed to run locally. Lists concrete packages, never `inline-core[...]` self-refs.
 all = [
   "torch>=2.2",
-  "diffusers>=0.36",
+  "diffusers>=0.39",
   "transformers>=4.44",
   "accelerate>=0.30",
   "safetensors>=0.4",

inline_core-1.2.5/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.3 → inline_core-1.2.5}/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.3 → inline_core-1.2.5}/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.3 → inline_core-1.2.5}/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.5/src/inline_core/models/checkpoint.py

@@ -0,0 +1,87 @@
+"""Read a safetensors checkpoint tensor by tensor, without mapping the whole file.
+
+``safetensors.safe_open`` maps the entire file at once, and Linux refuses a mapping larger than
+physical RAM when there is no swap (the default heuristic overcommit mode). A 26GB Krea 2 checkpoint
+is therefore unreadable on a 16GB machine - it fails with ``Cannot allocate memory`` before any GPU
+work, no matter how small the model would be once quantized.
+
+Reading each tensor's byte range instead keeps peak host RAM at one tensor and works for any file
+size. Only used for the big single-file checkpoints; small files (LoRAs, VAEs) still go through
+safetensors directly.
+"""
+
+from __future__ import annotations
+
+import json
+import struct
+from pathlib import Path
+from typing import Any
+
+from ..errors import ComponentError
+
+#: safetensors dtype names -> torch dtypes, resolved lazily so importing this module stays cheap.
+_DTYPE_NAMES = {
+    "F64": "float64",
+    "F32": "float32",
+    "F16": "float16",
+    "BF16": "bfloat16",
+    "I64": "int64",
+    "I32": "int32",
+    "I16": "int16",
+    "I8": "int8",
+    "U8": "uint8",
+    "BOOL": "bool",
+}
+
+
+class CheckpointReader:
+    """Random access to one safetensors file, one tensor at a time."""
+
+    def __init__(self, path: str | Path) -> None:
+        self._path = Path(path)
+        try:
+            size = self._path.stat().st_size
+            with self._path.open("rb") as handle:
+                (header_len,) = struct.unpack("<Q", handle.read(8))
+                # Garbage bytes decode to an enormous length; reject it rather than trying to
+                # allocate it, which would raise MemoryError instead of a usable message.
+                if not 0 < header_len <= size - 8:
+                    raise ValueError("header length is not plausible for this file")
+                header = json.loads(handle.read(header_len))
+            if not isinstance(header, dict):
+                raise ValueError("header is not a JSON object")
+        except (OSError, ValueError, struct.error) as error:
+            raise ComponentError(f"Could not read checkpoint {self._path.name}: {error}") from error
+        self._start = 8 + header_len
+        self._index: dict[str, Any] = {k: v for k, v in header.items() if k != "__metadata__"}
+        self.metadata: dict[str, Any] = header.get("__metadata__") or {}
+
+    def keys(self) -> list[str]:
+        return list(self._index)
+
+    def get_tensor(self, key: str, device: str | None = None) -> Any:
+        """One tensor, read straight from its byte range into a fresh buffer."""
+        import torch
+
+        entry = self._index[key]
+        dtype_name = _DTYPE_NAMES.get(entry["dtype"])
+        if dtype_name is None:
+            raise ComponentError(
+                f"Checkpoint {self._path.name} uses the unsupported dtype {entry['dtype']!r} "
+                f"for {key!r}."
+            )
+        start, end = entry["data_offsets"]
+        buffer = bytearray(end - start)
+        with self._path.open("rb") as handle:
+            handle.seek(self._start + start)
+            if handle.readinto(buffer) != len(buffer):
+                raise ComponentError(f"Checkpoint {self._path.name} is truncated at {key!r}.")
+        tensor = torch.frombuffer(buffer, dtype=getattr(torch, dtype_name))
+        tensor = tensor.reshape(entry["shape"]) if entry["shape"] else tensor.reshape(())
+        return tensor.to(device) if device else tensor
+
+    def __enter__(self) -> CheckpointReader:
+        return self
+
+    def __exit__(self, *_exc: object) -> None:
+        return None

inline_core-1.2.5/src/inline_core/models/krea2/__init__.py

@@ -0,0 +1,9 @@
+"""The Krea 2 (Krea AI) runtime: a diffusers-backed text-to-image / img2img runner.
+
+Optional subpackage. `server.bootstrap` imports `register_krea2` best-effort, so a core install
+without the ``runtime`` extra (torch + diffusers) still boots and serves the source nodes.
+"""
+
+from .runner import register_krea2
+
+__all__ = ["register_krea2"]

inline_core-1.2.5/src/inline_core/models/krea2/convert.py

@@ -0,0 +1,138 @@
+"""Rename a Krea 2 checkpoint from the reference (ComfyUI) layout to the diffusers one.
+
+``Krea2Transformer2DModel`` has no ``from_single_file``, so the single ``.safetensors`` under
+``diffusion_models/`` cannot be loaded directly. The two layouts describe the same 430 tensors under
+different names, so a pure rename plus one reshape is the whole conversion.
+
+The same rename maps a LoRA's module path (``module_alias``), which matters because the two Krea 2
+LoRA conventions disagree: the official Comfy-Org style LoRAs already use diffusers names, while
+ostris' training adapter uses reference ones.
+
+Torch-free apart from the tensors handed in, so it is cheap to import and easy to test.
+"""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Mapping
+from typing import Any
+
+from ...errors import ComponentError
+
+#: Root-anchored renames, longest-first so ``txtmlp.0.scale`` wins over the ``txtmlp.1.`` prefix.
+_ROOT: tuple[tuple[str, str], ...] = (
+    ("txtmlp.0.scale", "txt_in.norm.weight"),
+    ("txtmlp.1.", "txt_in.linear_1."),
+    ("txtmlp.3.", "txt_in.linear_2."),
+    ("tmlp.0.", "time_embed.linear_1."),
+    ("tmlp.2.", "time_embed.linear_2."),
+    ("tproj.1.", "time_mod_proj."),
+    ("first.", "img_in."),
+    ("last.linear.", "final_layer.linear."),
+    ("last.norm.scale", "final_layer.norm.weight"),
+    ("last.modulation.lin", "final_layer.scale_shift_table"),
+    ("txtfusion.", "text_fusion."),
+    ("blocks.", "transformer_blocks."),
+)
+
+#: Renames applied anywhere in the path - they name parts of a block, which appears at three depths
+#: (the main stack and the two text-fusion stacks).
+_INNER: tuple[tuple[str, str], ...] = (
+    (".attn.qknorm.qnorm.scale", ".attn.norm_q.weight"),
+    (".attn.qknorm.knorm.scale", ".attn.norm_k.weight"),
+    (".attn.wq.", ".attn.to_q."),
+    (".attn.wk.", ".attn.to_k."),
+    (".attn.wv.", ".attn.to_v."),
+    (".attn.wo.", ".attn.to_out.0."),
+    (".attn.gate.", ".attn.to_gate."),
+    (".mlp.up.", ".ff.up."),
+    (".mlp.gate.", ".ff.gate."),
+    (".mlp.down.", ".ff.down."),
+    (".prenorm.scale", ".norm1.weight"),
+    (".postnorm.scale", ".norm2.weight"),
+)
+
+#: Suffixes ComfyUI's quantized builds add on top of the 430 reference tensors.
+_QUANT_MARKERS = (".weight_scale", ".weight_scale_2", ".comfy_quant")
+
+_MOD_LIN = re.compile(r"\.mod\.lin$")
+
+
+def convert_key(key: str) -> str:
+    """One reference key in the diffusers naming. An already-diffusers key passes through."""
+    for old, new in _ROOT:
+        if key.startswith(old):
+            key = new + key[len(old) :]
+            break
+    for old, new in _INNER:
+        key = key.replace(old, new)
+    return _MOD_LIN.sub(".scale_shift_table", key)
+
+
+def module_alias(stem: str) -> str | None:
+    """A LoRA module path in diffusers naming, or None when the rename leaves it unchanged.
+
+    ``lora.py`` has already stripped the checkpoint prefix, so ``stem`` looks like
+    ``blocks.0.attn.wq``. The trailing ``.weight`` makes the suffix-anchored rules above fire."""
+    converted = convert_key(f"{stem}.weight").removesuffix(".weight")
+    return converted if converted != stem else None
+
+
+def is_quantized_checkpoint(keys: Any) -> bool:
+    """Whether this is a ComfyUI fp8/int8/nvfp4 build rather than the bf16 one."""
+    return any(str(k).endswith(_QUANT_MARKERS) for k in keys)
+
+
+def check_loadable(keys: Any, expected: set[str]) -> None:
+    """Refuse a checkpoint that would only partially load, before any tensor is read.
+
+    Key-only so the streaming loader can validate up front; a partial load would otherwise leave
+    random-initialised layers and produce quietly wrong images instead of an error."""
+    keys = list(keys)
+    if is_quantized_checkpoint(keys):
+        raise ComponentError(
+            "This is a ComfyUI quantized Krea 2 build (fp8 / int8 / nvfp4), which only ComfyUI can "
+            "read. Use krea2_raw_bf16.safetensors or krea2_turbo_bf16.safetensors - smart memory "
+            "quantizes it for your GPU on load."
+        )
+    converted = {convert_key(k): k for k in keys}
+    unknown = sorted(original for name, original in converted.items() if name not in expected)
+    if unknown:
+        raise ComponentError(
+            f"Krea 2 checkpoint has {len(unknown)} unrecognised tensors (e.g. "
+            f"{', '.join(unknown[:3])}). It is probably a different model or a quantized build."
+        )
+    missing = sorted(expected - set(converted))
+    if missing:
+        raise ComponentError(
+            f"Krea 2 checkpoint is missing {len(missing)} tensors (e.g. "
+            f"{', '.join(missing[:3])}). The file looks truncated or is not a full checkpoint."
+        )
+
+
+def convert_state_dict(
+    state: Mapping[str, Any], shapes: Mapping[str, tuple[int, ...]]
+) -> dict[str, Any]:
+    """Rename ``state`` into ``shapes``' naming, reshaping the one entry whose layout differs
+    (``blocks.N.mod.lin`` is flat where ``scale_shift_table`` is 2-D)."""
+    check_loadable(state.keys(), set(shapes))
+    converted: dict[str, Any] = {}
+    for key, tensor in state.items():
+        name = convert_key(key)
+        want = tuple(shapes[name])
+        if tuple(tensor.shape) != want:
+            if tensor.numel() != _numel(want):
+                raise ComponentError(
+                    f"Krea 2 checkpoint tensor {key!r} is {tuple(tensor.shape)}, but the model "
+                    f"expects {want}. This file is not a Krea 2 transformer."
+                )
+            tensor = tensor.reshape(want)
+        converted[name] = tensor
+    return converted
+
+
+def _numel(shape: tuple[int, ...]) -> int:
+    total = 1
+    for dim in shape:
+        total *= dim
+    return total