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.
Files changed (151) hide show
  1. {inline_core-1.2.3 → inline_core-1.2.31}/PKG-INFO +1 -1
  2. {inline_core-1.2.3 → inline_core-1.2.31}/pyproject.toml +1 -1
  3. inline_core-1.2.31/src/inline_core/ffmpeg.py +30 -0
  4. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/graph/loader_runners.py +32 -1
  5. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/graph/primitives.py +16 -1
  6. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/graph/schema.py +1 -0
  7. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/models/loaders.py +47 -12
  8. inline_core-1.2.31/src/inline_core/models/lora.py +140 -0
  9. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/models/zimage/runner.py +50 -10
  10. inline_core-1.2.31/tests/test_lora.py +218 -0
  11. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_studio_rpc.py +1 -0
  12. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_zimage_runner.py +10 -8
  13. {inline_core-1.2.3 → inline_core-1.2.31}/uv.lock +1 -1
  14. {inline_core-1.2.3 → inline_core-1.2.31}/.gitignore +0 -0
  15. {inline_core-1.2.3 → inline_core-1.2.31}/.python-version +0 -0
  16. {inline_core-1.2.3 → inline_core-1.2.31}/CLAUDE.md +0 -0
  17. {inline_core-1.2.3 → inline_core-1.2.31}/README.md +0 -0
  18. {inline_core-1.2.3 → inline_core-1.2.31}/main.py +0 -0
  19. {inline_core-1.2.3 → inline_core-1.2.31}/scripts/reference.py +0 -0
  20. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/__init__.py +0 -0
  21. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/components/__init__.py +0 -0
  22. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/components/conditioning.py +0 -0
  23. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/components/interfaces.py +0 -0
  24. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/config.py +0 -0
  25. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/device/__init__.py +0 -0
  26. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/device/auto.py +0 -0
  27. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/device/detect.py +0 -0
  28. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/device/memory.py +0 -0
  29. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/device/policy.py +0 -0
  30. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/device/types.py +0 -0
  31. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/errors.py +0 -0
  32. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/__init__.py +0 -0
  33. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/api.py +0 -0
  34. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/constraints.py +0 -0
  35. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/fetch.py +0 -0
  36. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/handlers.py +0 -0
  37. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/importer.py +0 -0
  38. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/install.py +0 -0
  39. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/loader.py +0 -0
  40. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/manifest.py +0 -0
  41. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/models.py +0 -0
  42. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/paths.py +0 -0
  43. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/resolve.py +0 -0
  44. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/scanner.py +0 -0
  45. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/state.py +0 -0
  46. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/extensions/tools.py +0 -0
  47. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/graph/__init__.py +0 -0
  48. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/graph/cache.py +0 -0
  49. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/graph/descriptor.py +0 -0
  50. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/graph/executor.py +0 -0
  51. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/graph/registry.py +0 -0
  52. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/graph/runners.py +0 -0
  53. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/graph/topo.py +0 -0
  54. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/graph/validate.py +0 -0
  55. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/media.py +0 -0
  56. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/models/__init__.py +0 -0
  57. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/models/catalog.py +0 -0
  58. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/models/requirements.py +0 -0
  59. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/models/sampling.py +0 -0
  60. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/models/zimage/__init__.py +0 -0
  61. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/models/zimage/primitives.py +0 -0
  62. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/models/zimage/provider.py +0 -0
  63. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/models/zimage/requirements.py +0 -0
  64. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/parallel/__init__.py +0 -0
  65. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/parallel/config.py +0 -0
  66. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/parallel/group.py +0 -0
  67. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/parallel/launch.py +0 -0
  68. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/parallel/protocol.py +0 -0
  69. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/parallel/registry.py +0 -0
  70. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/parallel/worker.py +0 -0
  71. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/runtime/__init__.py +0 -0
  72. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/runtime/context.py +0 -0
  73. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/runtime/file_store.py +0 -0
  74. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/runtime/progress.py +0 -0
  75. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/runtime/run.py +0 -0
  76. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/runtime/store.py +0 -0
  77. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/sampling/__init__.py +0 -0
  78. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/sampling/batch.py +0 -0
  79. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/server/__init__.py +0 -0
  80. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/server/__main__.py +0 -0
  81. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/server/app.py +0 -0
  82. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/server/assets.py +0 -0
  83. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/server/bootstrap.py +0 -0
  84. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/server/frontend.py +0 -0
  85. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/server/manager.py +0 -0
  86. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/server/rpc.py +0 -0
  87. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/server/run_store.py +0 -0
  88. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/server/serialize.py +0 -0
  89. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/__init__.py +0 -0
  90. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/assets.py +0 -0
  91. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/config.py +0 -0
  92. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/fal.py +0 -0
  93. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/frames.py +0 -0
  94. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/generation.py +0 -0
  95. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/graph_build.py +0 -0
  96. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/handlers.py +0 -0
  97. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/models.py +0 -0
  98. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/moodboard.py +0 -0
  99. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/peaks.py +0 -0
  100. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/schema.py +0 -0
  101. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/store.py +0 -0
  102. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/timeline/__init__.py +0 -0
  103. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/timeline/compose.py +0 -0
  104. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/timeline/ffmpeg.py +0 -0
  105. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/timeline/render.py +0 -0
  106. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/studio/timeline/resolve.py +0 -0
  107. {inline_core-1.2.3 → inline_core-1.2.31}/src/inline_core/takes.py +0 -0
  108. {inline_core-1.2.3 → inline_core-1.2.31}/tests/helpers.py +0 -0
  109. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_cache.py +0 -0
  110. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_catalog.py +0 -0
  111. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_config.py +0 -0
  112. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_device_detect.py +0 -0
  113. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_executor.py +0 -0
  114. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_extension_api.py +0 -0
  115. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_extension_install.py +0 -0
  116. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_extension_manifest.py +0 -0
  117. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_extension_resolve.py +0 -0
  118. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_extension_scanner.py +0 -0
  119. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_extension_spine.py +0 -0
  120. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_extension_state.py +0 -0
  121. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_file_store.py +0 -0
  122. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_frontend_serving.py +0 -0
  123. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_hidden_nodes.py +0 -0
  124. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_loader_runners.py +0 -0
  125. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_loaders.py +0 -0
  126. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_memory_policy.py +0 -0
  127. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_model_requirements.py +0 -0
  128. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_parallel_group.py +0 -0
  129. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_primitives.py +0 -0
  130. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_rpc_bridge.py +0 -0
  131. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_run_store.py +0 -0
  132. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_sampling.py +0 -0
  133. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_schema.py +0 -0
  134. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_server.py +0 -0
  135. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_studio_assets.py +0 -0
  136. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_studio_fal.py +0 -0
  137. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_studio_frames.py +0 -0
  138. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_studio_generation.py +0 -0
  139. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_studio_models.py +0 -0
  140. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_studio_moodboard.py +0 -0
  141. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_studio_peaks.py +0 -0
  142. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_studio_schema.py +0 -0
  143. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_studio_store.py +0 -0
  144. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_studio_timeline.py +0 -0
  145. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_take_bytes.py +0 -0
  146. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_topo.py +0 -0
  147. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_validate.py +0 -0
  148. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_xfuser_sampler.py +0 -0
  149. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_zimage_primitives.py +0 -0
  150. {inline_core-1.2.3 → inline_core-1.2.31}/tests/test_zimage_resolve.py +0 -0
  151. {inline_core-1.2.3 → inline_core-1.2.31}/webui.sh +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: inline-core
3
- Version: 1.2.3
3
+ Version: 1.2.31
4
4
  Summary: The generation engine behind Inline Studio.
5
5
  Requires-Python: >=3.11
6
6
  Requires-Dist: numpy>=1.26
@@ -1,7 +1,7 @@
1
1
  [project]
2
2
  # PyPI name; the import package is `inline_core` (src/inline_core).
3
3
  name = "inline-core"
4
- version = "1.2.3"
4
+ version = "1.2.31"
5
5
  description = "The generation engine behind Inline Studio."
6
6
  readme = "README.md"
7
7
  requires-python = ">=3.11"
@@ -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:
@@ -23,6 +23,7 @@ class PortKind(str, Enum):
23
23
  MODEL = "model"
24
24
  VAE = "vae"
25
25
  TEXT_ENCODER = "text-encoder"
26
+ LORA = "lora"
26
27
  CONDITIONING = "conditioning"
27
28
  LATENT = "latent"
28
29
 
@@ -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. The run manager
264
- # executes one run at a time; the lock guards each build.
265
- _CACHE: dict[tuple[str, str, str, str, str, str], Any] = {}
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, str, str, str, str, str], build: Callable[[], Any]) -> Any:
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 unload_components(keep_files: set[str] | None = None) -> None:
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 = [k for k in _CACHE if k[2] not in keep] # k = (arch, kind, file, dtype, quant, dev)
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
- return _cached(
340
- (arch, "diffusion", file, _dtype_key(dtype), quant.value, _device_key(device)), build
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
- _PIPELINES: dict[tuple[str, str, str, bool, str], Any] = {}
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
- key = (source, vae, text, img2img, quant.value)
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. Called under ``_LOCK``."""
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
- keep_triple = (source, vae, text)
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] == keep_triple:
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``."""