inline-core 1.2.41__tar.gz → 1.2.51__tar.gz

{inline_core-1.2.41 → inline_core-1.2.51}/PKG-INFO

@@ -1,13 +1,14 @@
 Metadata-Version: 2.4
 Name: inline-core
-Version: 1.2.41
+Version: 1.2.51
 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'
@@ -55,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
@@ -102,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/  ...
 ```
 
@@ -123,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
 
@@ -132,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
@@ -215,17 +231,18 @@ offload machinery, `--vram-budget GB` / `INLINE_VRAM_BUDGET_GB` to cap the assum
 
 ## Run
 
-The easy path is `webui.sh`, which maps friendly flags onto the engine's `INLINE_*` env knobs:
+The easy path is `webui.sh` (macOS/Linux) or `webui.bat` (Windows), which maps friendly flags onto the
+engine's `INLINE_*` env knobs. Windows must use `webui.bat` - `webui.sh` is a bash script:
 
 ```
-./webui.sh                            # loopback, port 8848
+./webui.sh                            # loopback, port 8848            (Windows: .\webui.bat)
 ./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
+`./webui.sh --help` (or `.\webui.bat --help`) lists every flag (networking, multi-GPU, device/memory profile, paths). The same
 flags are available on the Python entrypoint - `python main.py --help` - which is the dev path (it also
 takes `--front-end-root DIR` to serve a local SPA build). Or run the server module directly:
 

{inline_core-1.2.41 → inline_core-1.2.51}/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
@@ -165,17 +180,18 @@ offload machinery, `--vram-budget GB` / `INLINE_VRAM_BUDGET_GB` to cap the assum
 
 ## Run
 
-The easy path is `webui.sh`, which maps friendly flags onto the engine's `INLINE_*` env knobs:
+The easy path is `webui.sh` (macOS/Linux) or `webui.bat` (Windows), which maps friendly flags onto the
+engine's `INLINE_*` env knobs. Windows must use `webui.bat` - `webui.sh` is a bash script:
 
 ```
-./webui.sh                            # loopback, port 8848
+./webui.sh                            # loopback, port 8848            (Windows: .\webui.bat)
 ./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
+`./webui.sh --help` (or `.\webui.bat --help`) lists every flag (networking, multi-GPU, device/memory profile, paths). The same
 flags are available on the Python entrypoint - `python main.py --help` - which is the dev path (it also
 takes `--front-end-root DIR` to serve a local SPA build). Or run the server module directly:
 

{inline_core-1.2.41 → inline_core-1.2.51}/pyproject.toml

@@ -1,9 +1,10 @@
 [project]
 # PyPI name; the import package is `inline_core` (src/inline_core).
 name = "inline-core"
-version = "1.2.41"
+version = "1.2.51"
 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",
@@ -62,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.51/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.51/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.51/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

inline_core-1.2.51/src/inline_core/models/krea2/img2img.py

@@ -0,0 +1,89 @@
+"""Krea 2 img2img: start the schedule partway, from an encoded input image.
+
+diffusers ships no ``Krea2Img2ImgPipeline``, but ``Krea2Pipeline.__call__`` already accepts
+pre-noised ``latents`` and an explicit ``sigmas`` schedule, so img2img needs no fork of the denoise
+loop - only the right latents and the right tail of the schedule.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+
+def img2img_kwargs(
+    pipe: Any,
+    *,
+    image: Any,
+    strength: float,
+    steps: int,
+    width: int,
+    height: int,
+    generator: Any,
+    device: str,
+) -> dict[str, Any]:
+    """Start the schedule partway, from the input image, using the pipeline's own ``latents`` and
+    ``sigmas`` inputs - diffusers ships no ``Krea2Img2ImgPipeline`` and forking its denoise loop
+    would be a maintenance liability.
+
+    The sigma tail is handed over **unshifted**: ``set_timesteps`` applies the same monotone
+    exponential shift the pipeline would, so a shifted subset and a subset of the shifted schedule
+    are the same thing. The latents are noised at the *shifted* start sigma so the noise level the
+    model sees matches the timestep it is given."""
+    import numpy as np
+    from diffusers.pipelines.krea2.pipeline_krea2 import calculate_shift
+
+    raw = np.linspace(1.0, 1 / steps, steps)
+    start = min(steps - 1, max(0, int(round((1.0 - strength) * steps))))
+
+    vae_scale = pipe.vae_scale_factor * pipe.patch_size
+    grid_h, grid_w = height // vae_scale, width // vae_scale
+    if pipe.config.is_distilled:
+        mu = 1.15
+    else:
+        config = pipe.scheduler.config
+        mu = calculate_shift(
+            grid_h * grid_w,
+            config.get("base_image_seq_len", 256),
+            config.get("max_image_seq_len", 6400),
+            config.get("base_shift", 0.5),
+            config.get("max_shift", 1.15),
+        )
+    schedule = pipe.scheduler.__class__.from_config(pipe.scheduler.config)
+    schedule.set_timesteps(sigmas=raw, mu=mu, device=device)
+    sigma = schedule.sigmas[start].to(device)
+
+    latents = _encode_image(pipe, image, width, height, device, generator)
+    noise = torch.randn(
+        latents.shape, generator=generator, device=latents.device, dtype=latents.dtype
+    )
+    return {"latents": (1.0 - sigma) * latents + sigma * noise, "sigmas": raw[start:].tolist()}
+
+
+def _encode_image(
+    pipe: Any, image: Any, width: int, height: int, device: str, generator: Any
+) -> Any:
+    """The input image as packed, normalized Krea 2 latents. The Qwen-Image VAE is a video codec, so
+    the pixels carry a length-1 temporal axis. The VAE sample is drawn from the run's seeded
+    generator, not the global RNG, so the same seed reproduces the same img2img result."""
+    import numpy as np
+
+    resized = image.convert("RGB").resize((width, height))
+    array = np.asarray(resized, dtype="float32") / 127.5 - 1.0
+    dtype = pipe.vae.dtype
+    pixels = torch.from_numpy(array).permute(2, 0, 1)[None, :, None].to(device, dtype)
+
+    with torch.no_grad():
+        latents = pipe.vae.encode(pixels).latent_dist.sample(generator=generator)
+    mean, std = _latent_stats(pipe.vae, latents)
+    latents = ((latents - mean) / std).squeeze(2)  # drop the temporal axis
+    batch, channels, latent_h, latent_w = latents.shape
+    return pipe._pack_latents(latents, batch, channels, latent_h, latent_w)
+
+
+def _latent_stats(vae: Any, latents: Any) -> tuple[Any, Any]:
+    shape = (1, vae.config.z_dim, 1, 1, 1)
+    mean = torch.tensor(vae.config.latents_mean, device=latents.device, dtype=latents.dtype)
+    std = torch.tensor(vae.config.latents_std, device=latents.device, dtype=latents.dtype)
+    return mean.view(shape), std.view(shape)

inline_core-1.2.51/src/inline_core/models/krea2/provider.py

@@ -0,0 +1,57 @@
+"""Krea 2's answer to "what do I need on disk" - one provider per node variant."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+from ...config import models_dir
+from ..requirements import ModelComponent
+from .requirements import (
+    footprint_bytes,
+    krea2_requirements,
+    resolve_diffusion,
+    resolve_text_encoder,
+    resolve_vae,
+)
+
+
+class Krea2Provider:
+    """Requirements + fit estimate for one Krea 2 node (``turbo`` or ``raw``)."""
+
+    def __init__(self, variant: str) -> None:
+        self._variant = variant
+
+    def components(self, params: dict[str, object] | None = None) -> list[ModelComponent]:
+        return krea2_requirements(self._variant, params)
+
+    def download_target(self, component: ModelComponent) -> Path:
+        return models_dir() / component.category
+
+    def estimate(self, policy: Any) -> dict[str, Any] | None:
+        """Whether the model will fit this machine, so the popup can warn before a 26GB load.
+        ``None`` whenever it can't be sized - a wrong estimate is worse than none."""
+        if policy is None:
+            return None
+        try:
+            from ...device.policy import ModelFootprint
+        except ImportError:
+            return None
+        footprint = ModelFootprint(
+            **footprint_bytes(
+                resolve_diffusion(self._variant), resolve_vae(None), resolve_text_encoder(None)
+            )
+        )
+        fit = policy.estimate_fit(footprint)  # pure - never mutates the shared policy
+        if fit is None:
+            return None
+        soft = not fit.fits or fit.plan in ("int8", "offload")
+        return {
+            "plan": fit.plan,
+            "fits": fit.fits,
+            "requiredVramMb": int(fit.required_vram_gb * 1024),
+            "totalVramMb": int(fit.total_vram_gb * 1024) if fit.total_vram_gb else None,
+            "freeVramMb": policy.free_vram_mb(),
+            "freeRamMb": policy.free_ram_mb(),
+            "warning": fit.note if soft else None,
+        }