inline-core 1.2.3__tar.gz → 1.2.31__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {inline_core-1.2.3 → inline_core-1.2.31}/PKG-INFO +1 -1
- {inline_core-1.2.3 → inline_core-1.2.31}/pyproject.toml +1 -1
- inline_core-1.2.31/src/inline_core/ffmpeg.py +30 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/graph/loader_runners.py +32 -1
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/graph/primitives.py +16 -1
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/graph/schema.py +1 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/models/loaders.py +47 -12
- inline_core-1.2.31/src/inline_core/models/lora.py +140 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/models/zimage/runner.py +50 -10
- inline_core-1.2.31/tests/test_lora.py +218 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_studio_rpc.py +1 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_zimage_runner.py +10 -8
- {inline_core-1.2.3 → inline_core-1.2.31}/uv.lock +1 -1
- {inline_core-1.2.3 → inline_core-1.2.31}/.gitignore +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/.python-version +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/CLAUDE.md +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/README.md +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/main.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/scripts/reference.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/__init__.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/components/__init__.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/components/conditioning.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/components/interfaces.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/config.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/device/__init__.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/device/auto.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/device/detect.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/device/memory.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/device/policy.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/device/types.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/errors.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/__init__.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/api.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/constraints.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/fetch.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/handlers.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/importer.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/install.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/loader.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/manifest.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/models.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/paths.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/resolve.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/scanner.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/state.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/tools.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/graph/__init__.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/graph/cache.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/graph/descriptor.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/graph/executor.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/graph/registry.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/graph/runners.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/graph/topo.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/graph/validate.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/media.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/models/__init__.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/models/catalog.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/models/requirements.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/models/sampling.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/models/zimage/__init__.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/models/zimage/primitives.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/models/zimage/provider.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/models/zimage/requirements.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/parallel/__init__.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/parallel/config.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/parallel/group.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/parallel/launch.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/parallel/protocol.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/parallel/registry.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/parallel/worker.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/runtime/__init__.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/runtime/context.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/runtime/file_store.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/runtime/progress.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/runtime/run.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/runtime/store.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/sampling/__init__.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/sampling/batch.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/server/__init__.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/server/__main__.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/server/app.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/server/assets.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/server/bootstrap.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/server/frontend.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/server/manager.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/server/rpc.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/server/run_store.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/server/serialize.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/__init__.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/assets.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/config.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/fal.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/frames.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/generation.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/graph_build.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/handlers.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/models.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/moodboard.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/peaks.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/schema.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/store.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/timeline/__init__.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/timeline/compose.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/timeline/ffmpeg.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/timeline/render.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/timeline/resolve.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/takes.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/helpers.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_cache.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_catalog.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_config.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_device_detect.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_executor.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_extension_api.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_extension_install.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_extension_manifest.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_extension_resolve.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_extension_scanner.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_extension_spine.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_extension_state.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_file_store.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_frontend_serving.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_hidden_nodes.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_loader_runners.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_loaders.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_memory_policy.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_model_requirements.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_parallel_group.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_primitives.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_rpc_bridge.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_run_store.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_sampling.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_schema.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_server.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_studio_assets.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_studio_fal.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_studio_frames.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_studio_generation.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_studio_models.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_studio_moodboard.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_studio_peaks.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_studio_schema.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_studio_store.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_studio_timeline.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_take_bytes.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_topo.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_validate.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_xfuser_sampler.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_zimage_primitives.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_zimage_resolve.py +0 -0
- {inline_core-1.2.3 → inline_core-1.2.31}/webui.sh +0 -0
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
"""Locate the ffmpeg/ffprobe binaries. Prefers a bundled ``imageio-ffmpeg``, else PATH.
|
|
2
|
+
|
|
3
|
+
Lives at the top level rather than under ``studio/`` because both the timeline (studio) and the
|
|
4
|
+
take store (runtime) need it, and runtime must not import studio.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
import shutil
|
|
10
|
+
from functools import lru_cache
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
@lru_cache(maxsize=1)
|
|
14
|
+
def ffmpeg_exe() -> str | None:
|
|
15
|
+
try:
|
|
16
|
+
import imageio_ffmpeg
|
|
17
|
+
|
|
18
|
+
return imageio_ffmpeg.get_ffmpeg_exe()
|
|
19
|
+
except Exception: # noqa: BLE001
|
|
20
|
+
return shutil.which("ffmpeg")
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
@lru_cache(maxsize=1)
|
|
24
|
+
def ffprobe_exe() -> str | None:
|
|
25
|
+
"""PATH only - imageio bundles ffmpeg alone, so probing degrades gracefully when absent."""
|
|
26
|
+
return shutil.which("ffprobe")
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def ffmpeg_available() -> bool:
|
|
30
|
+
return ffmpeg_exe() is not None
|
|
@@ -20,7 +20,7 @@ from typing import TYPE_CHECKING, Any
|
|
|
20
20
|
|
|
21
21
|
from ..config import models_dir
|
|
22
22
|
from ..errors import ComponentError
|
|
23
|
-
from .primitives import LOAD_DIFFUSION_MODEL, LOAD_TEXT_ENCODER, LOAD_VAE
|
|
23
|
+
from .primitives import LOAD_DIFFUSION_MODEL, LOAD_LORA, LOAD_TEXT_ENCODER, LOAD_VAE
|
|
24
24
|
from .runners import NodeResult, NodeRunner
|
|
25
25
|
from .schema import Node
|
|
26
26
|
|
|
@@ -44,6 +44,16 @@ class ComponentRef:
|
|
|
44
44
|
file: str
|
|
45
45
|
|
|
46
46
|
|
|
47
|
+
@dataclass(frozen=True)
|
|
48
|
+
class LoraRef:
|
|
49
|
+
"""One LoRA in a stack: its absolute file path and blend strength. A ``load/lora`` node emits a
|
|
50
|
+
tuple of these (its own ref appended to any upstream stack); the model runner fuses them into
|
|
51
|
+
the diffusion transformer in order. Frozen + hashable so it can key the loader cache."""
|
|
52
|
+
|
|
53
|
+
file: str
|
|
54
|
+
strength: float
|
|
55
|
+
|
|
56
|
+
|
|
47
57
|
def _resolve_file(category: str, chosen: str) -> Path:
|
|
48
58
|
"""The single weight file a Load node points at: the explicit dropdown pick, else the first
|
|
49
59
|
weight file in ``models/<category>/`` (mirrors the model node's "auto"). Raises if none."""
|
|
@@ -81,6 +91,26 @@ class LoadComponentRunner(NodeRunner):
|
|
|
81
91
|
return NodeResult(outputs={self._output: ref})
|
|
82
92
|
|
|
83
93
|
|
|
94
|
+
class LoadLoraRunner(NodeRunner):
|
|
95
|
+
"""Resolve this node's ``file``/``strength`` into a ``LoraRef`` and append it to any upstream
|
|
96
|
+
stack on the ``lora`` input - so chaining ``load/lora`` nodes stacks them in wiring order. The
|
|
97
|
+
stack rides its own ``lora`` edge into the model runner, which fuses it into the transformer."""
|
|
98
|
+
|
|
99
|
+
produces_takes = False
|
|
100
|
+
|
|
101
|
+
def run(self, node: Node, inputs: dict[str, list[Any]], ctx: ExecutionContext) -> NodeResult:
|
|
102
|
+
upstream = _first(inputs.get("lora")) or ()
|
|
103
|
+
file = _resolve_file("loras", str(node.params.get("file", "")))
|
|
104
|
+
strength = float(node.params.get("strength", 1.0))
|
|
105
|
+
stack = (*upstream, LoraRef(file=str(file), strength=strength))
|
|
106
|
+
return NodeResult(outputs={"lora": stack})
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
def _first(values: list[Any] | None) -> Any:
|
|
110
|
+
"""The first wired value on a port, or None (an optional input may be absent/unconnected)."""
|
|
111
|
+
return values[0] if values else None
|
|
112
|
+
|
|
113
|
+
|
|
84
114
|
def register_loaders(registry: Registry) -> None:
|
|
85
115
|
"""Register the ``load/*`` nodes **visible** (unhidden) with their runners, so they appear in
|
|
86
116
|
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:
|
|
|
98
128
|
kind="text_encoder", category="text_encoders", output_port="text_encoder"
|
|
99
129
|
),
|
|
100
130
|
)
|
|
131
|
+
registry.register(replace(LOAD_LORA, hidden=False), LoadLoraRunner())
|
|
@@ -47,6 +47,20 @@ LOAD_TEXT_ENCODER = NodeDescriptor(
|
|
|
47
47
|
outputs=(Port("text_encoder", "Text encoder", PortKind.TEXT_ENCODER),),
|
|
48
48
|
)
|
|
49
49
|
|
|
50
|
+
LOAD_LORA = NodeDescriptor(
|
|
51
|
+
type="load/lora",
|
|
52
|
+
title="Load LoRA",
|
|
53
|
+
category="Loaders",
|
|
54
|
+
icon="box",
|
|
55
|
+
# Optional upstream lora: chain load/lora -> load/lora to stack several in order.
|
|
56
|
+
inputs=(Port("lora", "LoRA", PortKind.LORA, required=False),),
|
|
57
|
+
params=(
|
|
58
|
+
ParamField("file", "LoRA", Widget.SELECT, "", options_from="loras"),
|
|
59
|
+
ParamField("strength", "Strength", Widget.NUMBER, 1.0, min=-2.0, max=2.0, step=0.05),
|
|
60
|
+
),
|
|
61
|
+
outputs=(Port("lora", "LoRA", PortKind.LORA),),
|
|
62
|
+
)
|
|
63
|
+
|
|
50
64
|
ENCODE_TEXT = NodeDescriptor(
|
|
51
65
|
type="encode/text",
|
|
52
66
|
title="Encode Text",
|
|
@@ -122,6 +136,7 @@ PRIMITIVES: tuple[NodeDescriptor, ...] = (
|
|
|
122
136
|
LOAD_DIFFUSION_MODEL,
|
|
123
137
|
LOAD_VAE,
|
|
124
138
|
LOAD_TEXT_ENCODER,
|
|
139
|
+
LOAD_LORA,
|
|
125
140
|
ENCODE_TEXT,
|
|
126
141
|
EMPTY_LATENT,
|
|
127
142
|
SAMPLE,
|
|
@@ -132,7 +147,7 @@ PRIMITIVES: tuple[NodeDescriptor, ...] = (
|
|
|
132
147
|
|
|
133
148
|
#: The loader primitives now have runners and are offered in the add-node menu - registered
|
|
134
149
|
#: (unhidden) with their runners by ``graph/loader_runners.py``, so they are skipped here.
|
|
135
|
-
_HAS_RUNNER = {LOAD_DIFFUSION_MODEL.type, LOAD_VAE.type, LOAD_TEXT_ENCODER.type}
|
|
150
|
+
_HAS_RUNNER = {LOAD_DIFFUSION_MODEL.type, LOAD_VAE.type, LOAD_TEXT_ENCODER.type, LOAD_LORA.type}
|
|
136
151
|
|
|
137
152
|
|
|
138
153
|
def register_primitives(registry: Registry) -> None:
|
|
@@ -30,12 +30,15 @@ from collections.abc import Callable
|
|
|
30
30
|
from dataclasses import dataclass
|
|
31
31
|
from pathlib import Path
|
|
32
32
|
from threading import Lock
|
|
33
|
-
from typing import Any
|
|
33
|
+
from typing import TYPE_CHECKING, Any
|
|
34
34
|
|
|
35
35
|
from ..config import data_dir
|
|
36
36
|
from ..device.policy import Quantization
|
|
37
37
|
from ..errors import ComponentError
|
|
38
38
|
|
|
39
|
+
if TYPE_CHECKING:
|
|
40
|
+
from ..graph.loader_runners import LoraRef
|
|
41
|
+
|
|
39
42
|
|
|
40
43
|
@dataclass(frozen=True)
|
|
41
44
|
class ArchSpec:
|
|
@@ -260,13 +263,15 @@ def _quantize_in_place(model: Any, quant: Quantization) -> None:
|
|
|
260
263
|
# --- component loaders (cached in-process) ------------------------------------------------------
|
|
261
264
|
|
|
262
265
|
# Keyed by (arch, kind, file, dtype, quant, device) so switching one file (e.g. a different VAE),
|
|
263
|
-
# the quantization, or the load device reuses the other already-loaded components.
|
|
264
|
-
#
|
|
265
|
-
|
|
266
|
+
# the quantization, or the load device reuses the other already-loaded components. A fused diffusion
|
|
267
|
+
# transformer appends its LoRA stack (see ``lora_cache_key``) to that base key, so the same file
|
|
268
|
+
# with a different stack is a distinct entry. The run manager executes one run at a time; the lock
|
|
269
|
+
# guards each build.
|
|
270
|
+
_CACHE: dict[tuple[str, ...], Any] = {}
|
|
266
271
|
_CACHE_LOCK = Lock()
|
|
267
272
|
|
|
268
273
|
|
|
269
|
-
def _cached(key: tuple[str,
|
|
274
|
+
def _cached(key: tuple[str, ...], build: Callable[[], Any]) -> Any:
|
|
270
275
|
with _CACHE_LOCK:
|
|
271
276
|
hit = _CACHE.get(key)
|
|
272
277
|
if hit is not None:
|
|
@@ -284,17 +289,36 @@ def _device_key(device: str | None) -> str:
|
|
|
284
289
|
return device or "cpu"
|
|
285
290
|
|
|
286
291
|
|
|
287
|
-
def
|
|
292
|
+
def lora_cache_key(loras: tuple[LoraRef, ...]) -> tuple[str, ...]:
|
|
293
|
+
"""The cache-key suffix a LoRA stack contributes to a fused diffusion transformer. Order- and
|
|
294
|
+
strength-sensitive (fusing is not commutative), keyed by file+strength; an empty stack adds
|
|
295
|
+
nothing, so an un-LoRA'd load keys exactly as it did before LoRA existed."""
|
|
296
|
+
return tuple(f"{lora.file}@{lora.strength}" for lora in loras)
|
|
297
|
+
|
|
298
|
+
|
|
299
|
+
def unload_components(
|
|
300
|
+
keep_files: set[str] | None = None, keep_loras: tuple[str, ...] | None = None
|
|
301
|
+
) -> None:
|
|
288
302
|
"""Drop cached components whose source file is NOT in ``keep_files``, freeing their VRAM/RAM.
|
|
289
303
|
|
|
290
304
|
Called when switching checkpoints so a new model doesn't stack on top of the previous one (the
|
|
291
305
|
cache never evicted before, so a second distinct model roughly doubled resident memory). Only
|
|
292
306
|
drops references + empties the CUDA cache - it does not move weights to CPU RAM (that would just
|
|
293
307
|
relocate the pressure on a RAM-tight box). The caller must drop any pipeline holding these
|
|
294
|
-
components first, or the references keep them alive.
|
|
308
|
+
components first, or the references keep them alive.
|
|
309
|
+
|
|
310
|
+
``keep_loras`` (the LoRA-stack suffix being loaded, from ``lora_cache_key``) additionally evicts
|
|
311
|
+
a kept file's diffusion transformer when it carries a *different* stack: fusing a LoRA into an
|
|
312
|
+
already-resident checkpoint would otherwise keep the unfused transformer AND the fused one - two
|
|
313
|
+
full-size models on the card. Left None, eviction is file-only (the pre-LoRA behaviour)."""
|
|
295
314
|
keep = keep_files or set()
|
|
296
315
|
with _CACHE_LOCK:
|
|
297
|
-
stale = [
|
|
316
|
+
stale = []
|
|
317
|
+
for k in _CACHE: # k = (arch, kind, file, dtype, quant, dev, *lora_suffix)
|
|
318
|
+
if k[2] not in keep:
|
|
319
|
+
stale.append(k)
|
|
320
|
+
elif keep_loras is not None and k[1] == "diffusion" and tuple(k[6:]) != keep_loras:
|
|
321
|
+
stale.append(k)
|
|
298
322
|
for k in stale:
|
|
299
323
|
comp = _CACHE.pop(k)
|
|
300
324
|
del comp
|
|
@@ -307,6 +331,7 @@ def load_diffusion(
|
|
|
307
331
|
dtype: Any,
|
|
308
332
|
quant: Quantization = Quantization.NONE,
|
|
309
333
|
device: str | None = None,
|
|
334
|
+
loras: tuple[LoraRef, ...] = (),
|
|
310
335
|
) -> Any:
|
|
311
336
|
"""The diffusion transformer from a single ``.safetensors``. diffusers converts the checkpoint
|
|
312
337
|
keys; the config comes from the bundled assets, so nothing is fetched at load time. ``quant``
|
|
@@ -314,7 +339,8 @@ def load_diffusion(
|
|
|
314
339
|
**straight to the GPU** from an mmap-backed checkpoint - the fp16 weights are never materialized
|
|
315
340
|
as an anonymous CPU copy (the host-RAM spike that OOM-killed the server), and for the int8 path
|
|
316
341
|
torchao quantizes on-device per tensor. ``None`` loads to CPU (the offload path, where
|
|
317
|
-
accelerate installs its hooks before placing).
|
|
342
|
+
accelerate installs its hooks before placing). ``loras`` are fused into the weights in order
|
|
343
|
+
(the stack is part of the cache key), so a fused transformer is itself the cached artifact."""
|
|
318
344
|
|
|
319
345
|
def build() -> Any:
|
|
320
346
|
from diffusers import ZImageTransformer2DModel
|
|
@@ -333,12 +359,20 @@ def load_diffusion(
|
|
|
333
359
|
# ``from_pretrained``), so the transformer would load at full size and the "int8" plan would
|
|
334
360
|
# blow the VRAM budget (a T4 OOMs mid-load). Quantize it explicitly with torchao after the
|
|
335
361
|
# load instead - the weights briefly sit full-size on the device, then halve in place.
|
|
362
|
+
# Fuse LoRAs BEFORE quantizing: the fuse adds a full-precision delta into each weight, which
|
|
363
|
+
# int8 can't accept in place, and int8-quantized weights aren't a plain tensor to add into.
|
|
364
|
+
if loras:
|
|
365
|
+
from .lora import fuse_loras
|
|
366
|
+
|
|
367
|
+
fuse_loras(model, loras)
|
|
336
368
|
_quantize_in_place(model, quant)
|
|
337
369
|
return model
|
|
338
370
|
|
|
339
|
-
|
|
340
|
-
|
|
371
|
+
key = (
|
|
372
|
+
arch, "diffusion", file, _dtype_key(dtype), quant.value, _device_key(device),
|
|
373
|
+
*lora_cache_key(loras),
|
|
341
374
|
)
|
|
375
|
+
return _cached(key, build)
|
|
342
376
|
|
|
343
377
|
|
|
344
378
|
def load_vae(arch: str, file: str, dtype: Any, device: str | None = None) -> Any:
|
|
@@ -431,6 +465,7 @@ def assemble_zimage_pipeline(
|
|
|
431
465
|
quant: Quantization = Quantization.NONE,
|
|
432
466
|
vae_dtype: Any = None,
|
|
433
467
|
device: str | None = None,
|
|
468
|
+
loras: tuple[LoraRef, ...] = (),
|
|
434
469
|
cancel_check: Callable[[], None] | None = None,
|
|
435
470
|
) -> Any:
|
|
436
471
|
"""Build a Z-Image pipeline from three local single files. Components are cached individually,
|
|
@@ -448,7 +483,7 @@ def assemble_zimage_pipeline(
|
|
|
448
483
|
from diffusers import ZImageImg2ImgPipeline, ZImagePipeline
|
|
449
484
|
|
|
450
485
|
arch = _ZIMAGE.key
|
|
451
|
-
transformer = load_diffusion(arch, diffusion_file, dtype, quant, device=device)
|
|
486
|
+
transformer = load_diffusion(arch, diffusion_file, dtype, quant, device=device, loras=loras)
|
|
452
487
|
_release_transient()
|
|
453
488
|
if cancel_check is not None:
|
|
454
489
|
cancel_check()
|
|
@@ -0,0 +1,140 @@
|
|
|
1
|
+
"""Fuse LoRA weights into an already-loaded model, in stack order.
|
|
2
|
+
|
|
3
|
+
Fusing (rather than keeping adapters live) is deliberate: ``loaders.py`` caches components by a key
|
|
4
|
+
that includes the stack, so a fused model *is* the cached artifact and there is no adapter state to
|
|
5
|
+
get out of sync with that key.
|
|
6
|
+
|
|
7
|
+
Handles both key conventions we see in the wild - ComfyUI/kohya (``lora_down``/``lora_up`` + an
|
|
8
|
+
optional ``alpha``) and diffusers/peft (``lora_A``/``lora_B``). Every LoRA key must resolve to a
|
|
9
|
+
module: a partial match silently applies some of the LoRA and no-ops the rest, which degrades output
|
|
10
|
+
without erroring, so an unmatched key is a hard failure.
|
|
11
|
+
|
|
12
|
+
Known gap: fully underscore-flattened kohya paths (``lora_unet_layers_0_attention_to_q``) are
|
|
13
|
+
ambiguous - the reverse mapping cannot tell a path separator from an underscore inside a name like
|
|
14
|
+
``to_q``. Those raise rather than mis-apply. Every Z-Image LoRA seen so far uses dotted paths.
|
|
15
|
+
"""
|
|
16
|
+
|
|
17
|
+
from __future__ import annotations
|
|
18
|
+
|
|
19
|
+
from typing import TYPE_CHECKING, Any
|
|
20
|
+
|
|
21
|
+
from ..errors import ComponentError
|
|
22
|
+
|
|
23
|
+
if TYPE_CHECKING:
|
|
24
|
+
from ..graph.loader_runners import LoraRef
|
|
25
|
+
|
|
26
|
+
_DOWN = ("lora_down.weight", "lora_A.weight", "lora_A.default.weight")
|
|
27
|
+
_UP = ("lora_up.weight", "lora_B.weight", "lora_B.default.weight")
|
|
28
|
+
# Prefixes checkpoints put in front of the module path; stripped when matching against the model.
|
|
29
|
+
_PREFIXES = ("diffusion_model.", "transformer.", "lora_unet_", "lora_te_", "base_model.model.")
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def fuse_loras(model: Any, loras: tuple[LoraRef, ...]) -> None:
|
|
33
|
+
"""Merge each LoRA into ``model``'s weights in order. No-op for an empty stack."""
|
|
34
|
+
for lora in loras:
|
|
35
|
+
_fuse_one(model, lora.file, lora.strength)
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
def _fuse_one(model: Any, path: str, strength: float) -> None:
|
|
39
|
+
import torch
|
|
40
|
+
from safetensors.torch import load_file
|
|
41
|
+
|
|
42
|
+
try:
|
|
43
|
+
state = load_file(path)
|
|
44
|
+
except Exception as exc: # noqa: BLE001
|
|
45
|
+
raise ComponentError(f"Could not read LoRA {path!r}: {exc}") from exc
|
|
46
|
+
|
|
47
|
+
modules = _linear_modules(model)
|
|
48
|
+
pairs, alphas = _group(state)
|
|
49
|
+
if not pairs:
|
|
50
|
+
raise ComponentError(f"LoRA {path!r} contains no recognisable lora_down/lora_up pairs.")
|
|
51
|
+
|
|
52
|
+
unmatched: list[str] = []
|
|
53
|
+
fused = 0
|
|
54
|
+
for stem, (down, up) in sorted(pairs.items()):
|
|
55
|
+
target = _match(stem, modules)
|
|
56
|
+
if target is None:
|
|
57
|
+
unmatched.append(stem)
|
|
58
|
+
continue
|
|
59
|
+
with torch.no_grad():
|
|
60
|
+
weight = target.weight
|
|
61
|
+
delta = _delta(up, down, weight)
|
|
62
|
+
scale = strength * _alpha_scale(alphas.get(stem), down.shape[0])
|
|
63
|
+
weight.add_((delta * scale).to(device=weight.device, dtype=weight.dtype))
|
|
64
|
+
fused += 1
|
|
65
|
+
|
|
66
|
+
if unmatched:
|
|
67
|
+
sample = ", ".join(unmatched[:3])
|
|
68
|
+
raise ComponentError(
|
|
69
|
+
f"LoRA {path!r} does not match this model: {len(unmatched)} of {len(pairs)} layers "
|
|
70
|
+
f"have no target (e.g. {sample}). It was probably trained for a different architecture."
|
|
71
|
+
)
|
|
72
|
+
if fused == 0:
|
|
73
|
+
raise ComponentError(f"LoRA {path!r} matched no layers in this model.")
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
def _delta(up: Any, down: Any, weight: Any) -> Any:
|
|
77
|
+
"""``up @ down``, shaped to the target weight. Conv LoRAs flatten the spatial dims."""
|
|
78
|
+
dtype = _fuse_dtype(up)
|
|
79
|
+
delta = up.to("cpu", dtype=dtype).flatten(1) @ down.to("cpu", dtype=dtype).flatten(1)
|
|
80
|
+
return delta.reshape(weight.shape)
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
def _fuse_dtype(tensor: Any) -> Any:
|
|
84
|
+
import torch
|
|
85
|
+
|
|
86
|
+
# fp16 matmuls of low-rank factors lose precision; fuse in fp32 and cast on the way in.
|
|
87
|
+
return torch.float32 if tensor.dtype in (torch.float16, torch.bfloat16) else tensor.dtype
|
|
88
|
+
|
|
89
|
+
|
|
90
|
+
def _alpha_scale(alpha: Any, rank: int) -> float:
|
|
91
|
+
"""kohya-style checkpoints scale by ``alpha / rank``; without an alpha the factors are raw."""
|
|
92
|
+
if alpha is None or rank == 0:
|
|
93
|
+
return 1.0
|
|
94
|
+
return float(alpha.item() if hasattr(alpha, "item") else alpha) / float(rank)
|
|
95
|
+
|
|
96
|
+
|
|
97
|
+
def _group(state: dict[str, Any]) -> tuple[dict[str, tuple[Any, Any]], dict[str, Any]]:
|
|
98
|
+
downs: dict[str, Any] = {}
|
|
99
|
+
ups: dict[str, Any] = {}
|
|
100
|
+
alphas: dict[str, Any] = {}
|
|
101
|
+
for key, value in state.items():
|
|
102
|
+
for suffix in _DOWN:
|
|
103
|
+
if key.endswith("." + suffix):
|
|
104
|
+
downs[key[: -len(suffix) - 1]] = value
|
|
105
|
+
for suffix in _UP:
|
|
106
|
+
if key.endswith("." + suffix):
|
|
107
|
+
ups[key[: -len(suffix) - 1]] = value
|
|
108
|
+
if key.endswith(".alpha"):
|
|
109
|
+
alphas[key[: -len(".alpha")]] = value
|
|
110
|
+
return {k: (downs[k], ups[k]) for k in downs if k in ups}, alphas
|
|
111
|
+
|
|
112
|
+
|
|
113
|
+
def _linear_modules(model: Any) -> dict[str, Any]:
|
|
114
|
+
import torch
|
|
115
|
+
|
|
116
|
+
return {
|
|
117
|
+
name: module
|
|
118
|
+
for name, module in model.named_modules()
|
|
119
|
+
if isinstance(module, torch.nn.Linear | torch.nn.Conv2d | torch.nn.Conv3d)
|
|
120
|
+
}
|
|
121
|
+
|
|
122
|
+
|
|
123
|
+
def _match(stem: str, modules: dict[str, Any]) -> Any:
|
|
124
|
+
"""Resolve a checkpoint's module path to a module, tolerating the usual prefixes and the
|
|
125
|
+
underscore-flattened paths kohya emits."""
|
|
126
|
+
for candidate in _candidates(stem):
|
|
127
|
+
hit = modules.get(candidate)
|
|
128
|
+
if hit is not None:
|
|
129
|
+
return hit
|
|
130
|
+
return None
|
|
131
|
+
|
|
132
|
+
|
|
133
|
+
def _candidates(stem: str) -> list[str]:
|
|
134
|
+
out = [stem]
|
|
135
|
+
for prefix in _PREFIXES:
|
|
136
|
+
if stem.startswith(prefix):
|
|
137
|
+
out.append(stem[len(prefix) :])
|
|
138
|
+
# kohya flattens the whole path with underscores; the dotted form is the model's own naming.
|
|
139
|
+
out.extend([c.replace("_", ".") for c in list(out) if "." not in c])
|
|
140
|
+
return out
|
|
@@ -36,7 +36,7 @@ from ...device.policy import (
|
|
|
36
36
|
from ...device.types import DeviceKind, DType
|
|
37
37
|
from ...errors import CancelledError, ComponentError
|
|
38
38
|
from ...graph.descriptor import NodeDescriptor, ParamField, Port, Widget
|
|
39
|
-
from ...graph.loader_runners import ComponentRef
|
|
39
|
+
from ...graph.loader_runners import ComponentRef, LoraRef
|
|
40
40
|
from ...graph.runners import NodeResult, NodeRunner
|
|
41
41
|
from ...graph.schema import Node, PortKind
|
|
42
42
|
from ...media import MediaKind
|
|
@@ -112,6 +112,8 @@ ZIMAGE = NodeDescriptor(
|
|
|
112
112
|
Port("model", "Diffusion model", PortKind.MODEL, required=False),
|
|
113
113
|
Port("vae", "VAE", PortKind.VAE, required=False),
|
|
114
114
|
Port("text_encoder", "Text encoder", PortKind.TEXT_ENCODER, required=False),
|
|
115
|
+
# Optional LoRA stack (a load/lora chain) fused into the transformer before sampling.
|
|
116
|
+
Port("lora", "LoRA", PortKind.LORA, required=False),
|
|
115
117
|
# Optional: wire an image to run img2img instead of text-to-image.
|
|
116
118
|
Port("image", "Image (img2img)", PortKind.IMAGE, required=False),
|
|
117
119
|
),
|
|
@@ -183,6 +185,7 @@ class ZImageRunner(NodeRunner):
|
|
|
183
185
|
model_ref = _component_ref(inputs, "model", "diffusion")
|
|
184
186
|
vae_ref = _component_ref(inputs, "vae", "vae")
|
|
185
187
|
te_ref = _component_ref(inputs, "text_encoder", "text_encoder")
|
|
188
|
+
loras = _lora_stack(inputs)
|
|
186
189
|
wired = {ref.kind for ref in (model_ref, vae_ref, te_ref) if ref is not None}
|
|
187
190
|
|
|
188
191
|
# No hidden downloads: a required component that is neither wired nor on disk fails fast,
|
|
@@ -241,6 +244,7 @@ class ZImageRunner(NodeRunner):
|
|
|
241
244
|
vae=vae_file,
|
|
242
245
|
text=te_file,
|
|
243
246
|
quant=self._policy.quantization(),
|
|
247
|
+
loras=loras,
|
|
244
248
|
cancel_check=lambda: _raise_if_cancelled(ctx),
|
|
245
249
|
)
|
|
246
250
|
except CancelledError:
|
|
@@ -361,6 +365,11 @@ class ZImageRunner(NodeRunner):
|
|
|
361
365
|
"scheduler": scheduler,
|
|
362
366
|
"seed": seed,
|
|
363
367
|
**({"strength": call["strength"]} if img2img else {}),
|
|
368
|
+
**(
|
|
369
|
+
{"loras": [{"file": lo.file, "strength": lo.strength} for lo in loras]}
|
|
370
|
+
if loras
|
|
371
|
+
else {}
|
|
372
|
+
),
|
|
364
373
|
},
|
|
365
374
|
)
|
|
366
375
|
return NodeResult(outputs={"image": take}, takes=[take])
|
|
@@ -372,7 +381,9 @@ class ZImageRunner(NodeRunner):
|
|
|
372
381
|
# but the run manager executes one run at a time (workers=1). Switching any file or the quantization
|
|
373
382
|
# rebuilds - but the loader core caches each component, so only the changed one re-reads from disk.
|
|
374
383
|
# The lock guards concurrent first-time builds.
|
|
375
|
-
|
|
384
|
+
# (source, vae, text, img2img, quant, *lora_suffix) - img2img is a bool and the lora suffix is
|
|
385
|
+
# variable-length, so the key is a str|bool tuple rather than a fixed shape.
|
|
386
|
+
_PIPELINES: dict[tuple[str | bool, ...], Any] = {}
|
|
376
387
|
_LOCK = Lock()
|
|
377
388
|
|
|
378
389
|
|
|
@@ -385,9 +396,11 @@ def _load_pipeline(
|
|
|
385
396
|
vae: str,
|
|
386
397
|
text: str,
|
|
387
398
|
quant: Quantization = Quantization.NONE,
|
|
399
|
+
loras: tuple[LoraRef, ...] = (),
|
|
388
400
|
cancel_check: Callable[[], None] | None = None,
|
|
389
401
|
) -> Any:
|
|
390
|
-
|
|
402
|
+
lora_key = loaders.lora_cache_key(loras)
|
|
403
|
+
key = (source, vae, text, img2img, quant.value, *lora_key)
|
|
391
404
|
with _LOCK:
|
|
392
405
|
cached = _PIPELINES.get(key)
|
|
393
406
|
if cached is not None:
|
|
@@ -408,9 +421,9 @@ def _load_pipeline(
|
|
|
408
421
|
# Free any *other* model still resident before loading this one, so switching checkpoints
|
|
409
422
|
# doesn't stack VRAM/RAM (the caches never evicted before). Keeps the current source's
|
|
410
423
|
# components - including a base pipeline reused below for img2img.
|
|
411
|
-
_evict_stale(source, vae, text)
|
|
424
|
+
_evict_stale(source, vae, text, lora_key)
|
|
412
425
|
# An img2img pipe can reuse the base pipe's already-placed weights (no second load).
|
|
413
|
-
base = _PIPELINES.get((source, vae, text, False, quant.value))
|
|
426
|
+
base = _PIPELINES.get((source, vae, text, False, quant.value, *lora_key))
|
|
414
427
|
if img2img and base is not None:
|
|
415
428
|
pipe = ZImageImg2ImgPipeline.from_pipe(base)
|
|
416
429
|
logger.info(
|
|
@@ -434,6 +447,7 @@ def _load_pipeline(
|
|
|
434
447
|
quant=quant,
|
|
435
448
|
vae_dtype=vae_dtype,
|
|
436
449
|
device=load_device,
|
|
450
|
+
loras=loras,
|
|
437
451
|
cancel_check=cancel_check,
|
|
438
452
|
)
|
|
439
453
|
logger.info(
|
|
@@ -486,6 +500,7 @@ def _build_pipeline(
|
|
|
486
500
|
quant: Quantization = Quantization.NONE,
|
|
487
501
|
vae_dtype: Any = None,
|
|
488
502
|
device: str | None = None,
|
|
503
|
+
loras: tuple[LoraRef, ...] = (),
|
|
489
504
|
cancel_check: Callable[[], None] | None = None,
|
|
490
505
|
) -> Any:
|
|
491
506
|
"""Build a Z-Image pipeline **offline** - never touching the network.
|
|
@@ -508,6 +523,11 @@ def _build_pipeline(
|
|
|
508
523
|
"single-file layout (diffusion_models/ + vae/ + text_encoders/) to quantize.",
|
|
509
524
|
quant.value,
|
|
510
525
|
)
|
|
526
|
+
if loras:
|
|
527
|
+
raise ComponentError(
|
|
528
|
+
"LoRAs need the single-file diffusion layout (a file in diffusion_models/); they "
|
|
529
|
+
"can't be fused into a whole-pipeline folder. Remove the LoRA or switch the model."
|
|
530
|
+
)
|
|
511
531
|
cls = ZImageImg2ImgPipeline if img2img else ZImagePipeline
|
|
512
532
|
return cls.from_pretrained(source, torch_dtype=dtype, local_files_only=True)
|
|
513
533
|
|
|
@@ -525,23 +545,29 @@ def _build_pipeline(
|
|
|
525
545
|
quant=quant,
|
|
526
546
|
vae_dtype=vae_dtype,
|
|
527
547
|
device=device,
|
|
548
|
+
loras=loras,
|
|
528
549
|
cancel_check=cancel_check,
|
|
529
550
|
)
|
|
530
551
|
|
|
531
552
|
|
|
532
|
-
def _evict_stale(source: str, vae: str, text: str) -> None:
|
|
553
|
+
def _evict_stale(source: str, vae: str, text: str, lora_key: tuple[str, ...]) -> None:
|
|
533
554
|
"""Free every *other* model's pipelines + components before loading a new checkpoint, so a
|
|
534
555
|
second distinct model doesn't stack on the first. Entries for the current ``(source, vae,
|
|
535
|
-
text)`` - e.g. a cached base pipeline reused for img2img - are kept.
|
|
556
|
+
text)`` - e.g. a cached base pipeline reused for img2img - are kept. ``lora_key`` (the stack
|
|
557
|
+
about to be fused) additionally drops a kept file's transformer if it carries a *different*
|
|
558
|
+
stack, so switching LoRAs on the same checkpoint doesn't keep both fused variants resident.
|
|
559
|
+
Called under ``_LOCK``."""
|
|
536
560
|
import gc
|
|
537
561
|
|
|
538
|
-
|
|
562
|
+
# Keep only the current (files, LoRA stack)'s pipelines - both img2img variants (k[3]) of the
|
|
563
|
+
# SAME lora suffix (k[5:]). A different stack of the same files must go, or its pipeline keeps the
|
|
564
|
+
# unfused transformer alive and a second (fused) one loads on top -> two transformers -> OOM.
|
|
539
565
|
for k in list(_PIPELINES):
|
|
540
|
-
if k[:3] ==
|
|
566
|
+
if k[:3] == (source, vae, text) and tuple(k[5:]) == lora_key:
|
|
541
567
|
continue
|
|
542
568
|
del _PIPELINES[k]
|
|
543
569
|
gc.collect()
|
|
544
|
-
loaders.unload_components(keep_files={source, vae, text})
|
|
570
|
+
loaders.unload_components(keep_files={source, vae, text}, keep_loras=lora_key)
|
|
545
571
|
_free_vram()
|
|
546
572
|
|
|
547
573
|
|
|
@@ -883,6 +909,20 @@ def _component_ref(inputs: dict[str, list[Any]], port: str, kind: str) -> Compon
|
|
|
883
909
|
raise ComponentError(f"Z-Image '{port}' input is not a loadable {kind} handle.")
|
|
884
910
|
|
|
885
911
|
|
|
912
|
+
def _lora_stack(inputs: dict[str, list[Any]]) -> tuple[LoraRef, ...]:
|
|
913
|
+
"""The LoRA stack wired into the optional ``lora`` port (a ``load/lora`` chain's output), or an
|
|
914
|
+
empty stack when unwired. The graph validator guarantees the port kind, so the value is a tuple
|
|
915
|
+
of ``LoraRef`` already in fuse order."""
|
|
916
|
+
stack = _first(inputs.get("lora"))
|
|
917
|
+
if stack is None:
|
|
918
|
+
return ()
|
|
919
|
+
if isinstance(stack, tuple):
|
|
920
|
+
items: tuple[Any, ...] = stack
|
|
921
|
+
if all(isinstance(item, LoraRef) for item in items):
|
|
922
|
+
return cast("tuple[LoraRef, ...]", items)
|
|
923
|
+
raise ComponentError("Z-Image 'lora' input is not a LoRA handle.")
|
|
924
|
+
|
|
925
|
+
|
|
886
926
|
def _path_or_none(path: Any) -> str:
|
|
887
927
|
"""A resolved component path as a string, or ``""`` when absent - the pipeline cache key and the
|
|
888
928
|
single-file check both want a plain string, not ``Path | None``."""
|