Inline Studio is a free, open-source app for AI filmmakers. Build a whole visual pipeline on a free-form node canvas, from moodboard to final cut, with local diffusion models (the built-in Inline Core engine) and hosted fal models. Every render is kept as a versioned, non-destructive take.
+
Inline Studio is a free, open-source app for AI filmmakers. Build a whole visual pipeline on a free-form node canvas, from moodboard to final cut, with local diffusion models (the built-in Inline Core engine) and hosted fal models. Train your own LoRAs on the same canvas. Every render is kept as a versioned, non-destructive take.
-
+
@@ -35,6 +35,7 @@ It runs as a **single process on one port**: the Inline Core engine (Python) ser
- **Chain frames into a generative pipeline** - wire one frame's output into the next frame's input. Refine a shot, feed it forward, regenerate the source, and everything downstream follows.
- **Video editing on the canvas** - the **Video Director node** is a timeline-in-a-node that assembles your rendered frames into a single cut, with layered audio (the videos' own audio plus your own music/VO), per-input and per-layer volume, an in-node preview to scrub, and high-res export; the **Trim Video/Audio node** lets you drop in a clip, drag the in/out handles over its filmstrip/waveform, and pass just the trimmed segment downstream.
- **Local generation, built in** - the Inline Core engine runs diffusion models on your own GPU. Z-Image Turbo from a single model file, no external server to set up.
+- **Train your own LoRAs** - the Trainer tab is a second canvas where the dataset, captioning, training run, and loss curve are all nodes. The finished LoRA drops into `models/loras/` and shows up in the LoRA loader node, ready to generate with. See [LoRA training](#lora-training).
- **API Nodes for hosted models** - run closed models right on the canvas with no GPU. Add a Generate node, pick a model, and bring your own provider key. See [API Nodes](#api-nodes).
- **Community extensions** - install custom nodes from a GitHub repo in one click, security-reviewed and dependency-isolated. Browse the [registry](https://github.com/inlineresearch/Inline-Registry) or [build your own](https://github.com/inlineresearch/Inline-Studio-Extension-Guide).
- **Free & open source (MIT)** - one process (Python + a browser); runs on macOS, Windows, and Linux.
@@ -228,6 +229,68 @@ The friendly `webui.sh` launcher (in `core/`) maps flags onto the engine's `INLI
`webui.sh` also has `--install` / `--extra NAME` to set up the venv. New to Inline Studio? The [Getting Started guide](https://inlinestudio.art/getting-started) walks you through your first render.
+## LoRA training
+
+Train a LoRA on your own images without leaving the app. The **Trainer** tab is a second canvas: wire up the nodes, press Start, and watch it run. When the run finishes, the `.safetensors` lands in `models/loras/`, where the LoRA loader node picks it up automatically, so you can generate with it over in the Studio tab straight away.
+
+
+
+### The graph
+
+Five nodes, wired left to right:
+
+- **Load Dataset** picks a training dataset and feeds it downstream. The node face stays a preview (thumbnails, image and caption counts); the images and captions themselves are edited in the side panel.
+- **Caption** runs a local captioner over the images that need one, with per-image progress. Captions stay editable afterwards, and a wired dataset overrides the node's own picker.
+- **Train LoRA** runs the job. Hyperparameters live behind the Adjust button, off the node face, so the node stays a status surface: a live step counter, the trainer's streaming logs, and a progress bar. The run control is a single chip that reads Start, Stop, or Resume depending on where the run is.
+- **Graph** plots the loss curve for whichever run is wired into it, with loss values on the y axis and the step range on the x axis.
+- **Resources** is a read-only readout of CPU, RAM, and VRAM as circular gauges. It takes no connections, and you can drop it on the Studio canvas too.
+
+
+Datasets, resuming, trigger words, and base model modes
+
+### Datasets and outputs
+
+The sidebar has two tabs. **Datasets** is where you create a dataset, give it a trigger word, add images (drag and drop from your file manager works), and edit captions. **Outputs** lists what training has produced: finished LoRAs with their rank, step count, and resolution, plus any run that stopped early, each with a Resume button.
+
+### Stop and resume
+
+Stopping a run flushes a checkpoint before the process exits, so Resume continues from the step it left off instead of starting over. A checkpoint holds the adapter weights, the optimiser state, the RNG state, and the step number, which is what makes a resumed run a continuation rather than a restart. Runs cut short by a crash or a server restart are recovered the same way and show up under Outputs ready to resume.
+
+### Trigger words
+
+A dataset's trigger word is prepended to every caption during training, so the model sees captions in the form `mytoken, a photo of ...`. Put the same token at the front of your prompt to pull the LoRA in. It is worth matching the phrasing of your captions too: if they all say "an oil painting of", a prompt written the same way will hit the trained style far more reliably than the trigger word alone.
+
+### Base model modes
+
+Z-Image Turbo is step-distilled, and training directly on a distilled model breaks the distillation down (turbo drift). There are two ways around it:
+
+- **Turbo + training adapter** fuses a de-distillation adapter into the base for the duration of training and drops it when the LoRA is saved, which preserves the 8-step speed. Put [ostris/zimage_turbo_training_adapter](https://huggingface.co/ostris/zimage_turbo_training_adapter) in `models/loras/`; any filename containing `adapter` is detected automatically, or point `INLINE_ZIMAGE_TRAIN_ADAPTER` at a specific file. Keep runs short, since the adapter slows the breakdown rather than preventing it.
+- **De-Turbo** trains without an adapter and needs no extra download.
+
+
+
+### Install
+
+The trainer's dependencies (PEFT, 8-bit Adam, the captioner) sit behind the `training` extra, so a normal install stays lean:
+
+```bash
+cd core
+./webui.sh --install --extra runtime --extra training
+```
+
+Weights are bring-your-own, the same as generation: nothing is downloaded behind your back. Training reuses the Z-Image files you already have in `models/diffusion_models/`, `models/vae/`, and `models/text_encoders/`. The captioner is the one exception, fetched once into the Hugging Face cache the first time you press Auto-caption.
+
+### Hardware
+
+Training needs a real GPU, and the resolution you train at drives peak VRAM far more than rank or batch size does. Two configurations have been run end to end:
+
+| GPU | Resolution | Notes |
+| --------------- | ---------- | -------------------------------------------------------------------- |
+| 16GB (Tesla T4) | 512px | Peaks around 13GB. 768 and 1024 both run out of memory on this card. |
+| 24GB (L4) | 1024px | Run with the Turbo training adapter fused in. |
+
+To keep the peak down, the VAE and text encoder are loaded first, used to cache latents and captions, then freed before the transformer loads, so the peak is the transformer on its own rather than all three resident at once. If you do hit an out-of-memory error, lower the training resolution before changing anything else.
+
## FAQ
### Is Inline Studio free?
@@ -242,10 +305,6 @@ Only for **local** generation. The built-in Inline Core engine renders on the GP
See [Two ways to generate](#two-ways-to-generate): local Z-Image on your own GPU, or hosted fal models. Adding a new local model is a Core change (a model runner), no UI release.
-### Does it use ComfyUI?
-
-No. Built-in generation is all Inline Core (local, on your own GPU) plus hosted fal models — nothing external to stand up. There's no ComfyUI connection, Generate tab, or workflow linking.
-
## Contributing
Inline Studio is early and moving fast, any issues, ideas, and pull requests are all welcome. If you're poking at the code, [CLAUDE.md](CLAUDE.md) is the engineering guide: it explains the architecture, the data model, and the conventions to follow.
@@ -256,6 +315,8 @@ Want to help by using it for real? Try the [creator task](task.md): build a shor
Inline Core's multi-GPU denoise builds on [**xDiT**](https://github.com/xdit-project/xDiT)'s PipeFusion and Ulysses parallelism.
+The LoRA trainer's approach to training on a step-distilled model follows [**ai-toolkit**](https://github.com/ostris/ai-toolkit) by ostris, and Turbo mode uses his [Z-Image Turbo training adapter](https://huggingface.co/ostris/zimage_turbo_training_adapter).
+
## Help shape Inline Studio
Are you an AI filmmaker who wants to help us make this better? We run a **paid trial feedback program**: use Inline Studio on real work, tell us what helps and what gets in your way, and get paid for your time.
diff --git a/core/pyproject.toml b/core/pyproject.toml
index db525b8..f1e4c8e 100644
--- a/core/pyproject.toml
+++ b/core/pyproject.toml
@@ -1,7 +1,7 @@
[project]
# PyPI name; the import package is `inline_core` (src/inline_core).
name = "inline-core"
-version = "1.2.31"
+version = "1.2.41"
description = "The generation engine behind Inline Studio."
readme = "README.md"
requires-python = ">=3.11"
@@ -39,6 +39,20 @@ parallel = [
"xfuser>=0.4",
"nvidia-ml-py>=12",
]
+# LoRA training (the Trainer tab): PEFT adapter training + local auto-caption + host/GPU telemetry.
+# Install alongside `runtime`: `.[runtime,training]`. Reuses runtime's torch/diffusers/accelerate.
+training = [
+ "peft>=0.11",
+ # 8-bit Adam keeps optimizer state small; no macOS wheels, so skip it there (AdamW fallback).
+ "bitsandbytes>=0.43; platform_system != 'Darwin'",
+ "Pillow>=10",
+ # Telemetry (CPU/RAM + per-GPU NVML) for the Trainer tab.
+ "psutil>=5.9",
+ "nvidia-ml-py>=12",
+ # Florence-2 auto-captioner is loaded via trust_remote_code and needs these.
+ "timm>=1.0",
+ "einops>=0.7",
+]
dev = [
"pytest>=8",
"ruff>=0.6",
diff --git a/core/src/inline_core/server/app.py b/core/src/inline_core/server/app.py
index 927abaf..e450392 100644
--- a/core/src/inline_core/server/app.py
+++ b/core/src/inline_core/server/app.py
@@ -29,6 +29,7 @@
from ..models.catalog import ModelCatalog
from ..models.requirements import RequirementsRegistry
from ..runtime.file_store import FileTakeStore
+from ..studio.system_stats import SystemStats
from .assets import AssetStore
from .manager import RunConflict, RunManager
from .rpc import EventBroadcaster, RpcRouter
@@ -139,13 +140,19 @@ def create_app(
manager = RunManager(registry, cache, policy, store=run_store, takes=FileTakeStore(takes_root))
rpc = rpc or RpcRouter()
events = events or EventBroadcaster()
+ # Host/GPU telemetry for the Trainer tab; only meaningful with the SPA (studio) backend wired.
+ stats = SystemStats(events) if studio_store is not None else None
@asynccontextmanager
async def lifespan(app: FastAPI): # noqa: ANN202
manager.bind_loop(asyncio.get_running_loop())
catalog.ensure_dirs()
catalog.scan()
+ if stats is not None:
+ stats.start()
yield
+ if stats is not None:
+ stats.stop()
manager.shutdown()
app = FastAPI(title="Inline Core", version="0.0.0", lifespan=lifespan)
@@ -298,6 +305,7 @@ async def studio_events(websocket: WebSocket) -> None:
from ..studio.handlers import register_studio_handlers
from ..studio.models import ModelDownloads
from ..studio.timeline.render import Timeline
+ from ..studio.training import Training
def core_models() -> dict[str, Any]:
return {
@@ -316,6 +324,7 @@ def core_status() -> dict[str, Any]:
generation=CoreGeneration(studio_store, manager, events),
fal_generation=FalGeneration(studio_store, events),
timeline=Timeline(studio_store, events),
+ training=Training(studio_store, events, on_output=catalog.rescan),
# Explicit model downloads write into models/; rescan so new files bump the registry.
# The policy lets the requirements popup show a memory fit estimate before a load;
# the requirements registry says which node types have models at all.
diff --git a/core/src/inline_core/server/rpc.py b/core/src/inline_core/server/rpc.py
index 36ee35c..8a0ad30 100644
--- a/core/src/inline_core/server/rpc.py
+++ b/core/src/inline_core/server/rpc.py
@@ -56,6 +56,11 @@ class EventBroadcaster:
def __init__(self) -> None:
self._subscribers: set[asyncio.Queue[dict[str, Any]]] = set()
+ @property
+ def subscriber_count(self) -> int:
+ """Open ``/events`` sockets - lets pollers (telemetry) skip work when nobody listens."""
+ return len(self._subscribers)
+
def add(self) -> asyncio.Queue[dict[str, Any]]:
queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
self._subscribers.add(queue)
diff --git a/core/src/inline_core/studio/handlers.py b/core/src/inline_core/studio/handlers.py
index 840b473..715d633 100644
--- a/core/src/inline_core/studio/handlers.py
+++ b/core/src/inline_core/studio/handlers.py
@@ -39,6 +39,7 @@ def register_studio_handlers(
generation: Any = None,
fal_generation: Any = None,
timeline: Any = None,
+ training: Any = None,
model_downloads: Any = None,
app_version: str = "1.0.0",
) -> None:
@@ -155,7 +156,9 @@ def delete_take(take_id: str) -> None:
reg("frames:deleteTake", delete_take)
# --- moodboard ------------------------------------------------------------------------------
- reg("moodboard:list", lambda: mb.list_board(conn()))
+ # `surface` defaults to the Studio moodboard so existing callers are unchanged; the Trainer tab
+ # passes "trainer" to get its own isolated canvas out of the same tables.
+ reg("moodboard:list", lambda surface=mb.STUDIO_SURFACE: mb.list_board(conn(), surface))
reg("moodboard:addAsset", lambda aid, x, y: mb.add_asset(conn(), aid, x, y))
reg("moodboard:addText", lambda x, y: mb.add_text(conn(), x, y))
reg("moodboard:addFrameFromAsset", lambda aid, x, y: mb.add_frame_from_asset(conn(), aid, x, y))
@@ -183,7 +186,18 @@ def delete_take(take_id: str) -> None:
reg("moodboard:setConnectorVolume", lambda cid, vol: mb.set_connector_volume(conn(), cid, vol))
reg(
"moodboard:replaceBoard",
- lambda items, connectors: mb.replace_board(conn(), items, connectors),
+ lambda items, connectors, surface=mb.STUDIO_SURFACE: mb.replace_board(
+ conn(), items, connectors, surface
+ ),
+ )
+ # Trainer-canvas nodes (plus the shared read-only resource node, which either canvas can host).
+ reg("moodboard:addTrainDataset", lambda x, y: mb.add_train_dataset(conn(), x, y))
+ reg("moodboard:addCaption", lambda x, y: mb.add_caption(conn(), x, y))
+ reg("moodboard:addTrainer", lambda x, y: mb.add_trainer(conn(), x, y))
+ reg("moodboard:addLossGraph", lambda x, y: mb.add_loss_graph(conn(), x, y))
+ reg(
+ "moodboard:addResource",
+ lambda x, y, surface=mb.STUDIO_SURFACE: mb.add_resource(conn(), x, y, surface),
)
# --- generation -----------------------------------------------------------------------------
@@ -207,6 +221,25 @@ def cancel_generation(frame_id: str | None = None) -> None:
reg("generation:cancel", cancel_generation)
reg("generation:resumePending", lambda: None)
+ # --- LoRA training (dataset CRUD + the training run subprocess) ------------------------------
+ if training is not None:
+ reg("training:listDatasets", lambda: training.list_datasets())
+ reg("training:createDataset", lambda inp: training.create_dataset(inp))
+ reg("training:listItems", lambda did: training.list_items(did))
+ reg("training:addItems", lambda did, aids: training.add_items(did, aids))
+ reg("training:removeItem", lambda iid: training.remove_item(iid))
+ reg("training:setCaption", lambda iid, cap: training.set_caption(iid, cap))
+ reg("training:autoCaption", lambda did, overwrite: training.auto_caption(did, overwrite))
+ reg("training:listRuns", lambda: training.list_runs())
+ reg("training:start", lambda did, hp: training.start(did, hp))
+ reg("training:resume", lambda rid: training.resume(rid))
+ reg("training:cancel", lambda rid: training.cancel(rid))
+ reg("training:status", lambda rid: training.status(rid))
+ else:
+ for ch in ("listDatasets", "createDataset", "listItems", "addItems", "removeItem",
+ "setCaption", "autoCaption", "listRuns", "start", "resume", "cancel", "status"):
+ reg(f"training:{ch}", not_wired("LoRA training"))
+
# --- fal settings (key stored server-side) --------------------------------------------------
reg("falSettings:status", store.fal_status)
reg("falSettings:setApiKey", store.set_fal_key)
diff --git a/core/src/inline_core/studio/moodboard.py b/core/src/inline_core/studio/moodboard.py
index 8973637..c3c63d1 100644
--- a/core/src/inline_core/studio/moodboard.py
+++ b/core/src/inline_core/studio/moodboard.py
@@ -17,10 +17,13 @@
from . import frames as fr
_ITEM_COLUMNS = (
- "id, project_id, type, asset_id, frame_id, parent_id, data, x, y, width, height, rotation, "
- "z_index, created_at, updated_at"
+ "id, project_id, surface, type, asset_id, frame_id, parent_id, data, x, y, width, height, "
+ "rotation, z_index, created_at, updated_at"
)
+#: The Studio moodboard; the Trainer tab's graph is the other surface (see schema `surface`).
+STUDIO_SURFACE = "studio"
+
_DEFAULT_SIZE = {"image": (320, 180), "video": (360, 203), "audio": (320, 80)}
_DEFAULT_TEXT = {
@@ -63,6 +66,7 @@ def _row_to_item(row: sqlite3.Row) -> dict[str, Any]:
return {
"id": row["id"],
"projectId": row["project_id"],
+ "surface": row["surface"],
"type": row["type"],
"assetId": row["asset_id"],
"frameId": row["frame_id"],
@@ -83,6 +87,7 @@ def _row_to_connector(row: sqlite3.Row) -> dict[str, Any]:
return {
"id": row["id"],
"projectId": row["project_id"],
+ "surface": row["surface"],
"fromItemId": row["from_item_id"],
"toItemId": row["to_item_id"],
"label": row["label"],
@@ -103,16 +108,28 @@ def list_items(conn: sqlite3.Connection) -> list[dict[str, Any]]:
return [_row_to_item(r) for r in rows]
-def list_board(conn: sqlite3.Connection) -> dict[str, Any]:
- items = [_row_to_item(r) for r in conn.execute("SELECT * FROM moodboard_items").fetchall()]
+def list_board(conn: sqlite3.Connection, surface: str = STUDIO_SURFACE) -> dict[str, Any]:
+ """One canvas's items + connectors. Surfaces are isolated: the Trainer graph never leaks into
+ the Studio moodboard (and vice versa)."""
+ items = [
+ _row_to_item(r)
+ for r in conn.execute(
+ "SELECT * FROM moodboard_items WHERE surface = ?", (surface,)
+ ).fetchall()
+ ]
connectors = [
- _row_to_connector(r) for r in conn.execute("SELECT * FROM moodboard_connectors").fetchall()
+ _row_to_connector(r)
+ for r in conn.execute(
+ "SELECT * FROM moodboard_connectors WHERE surface = ?", (surface,)
+ ).fetchall()
]
return {"items": items, "connectors": connectors}
-def _next_z(conn: sqlite3.Connection) -> int:
- row = conn.execute("SELECT MAX(z_index) AS z FROM moodboard_items").fetchone()
+def _next_z(conn: sqlite3.Connection, surface: str = STUDIO_SURFACE) -> int:
+ row = conn.execute(
+ "SELECT MAX(z_index) AS z FROM moodboard_items WHERE surface = ?", (surface,)
+ ).fetchone()
return (row["z"] or 0) + 1
@@ -128,11 +145,13 @@ def _insert_item(
asset_id: str | None = None,
frame_id: str | None = None,
z_index: int | None = None,
+ surface: str = STUDIO_SURFACE,
) -> dict[str, Any]:
now = _now()
item = {
"id": _uuid(),
"project_id": _project_id(conn),
+ "surface": surface,
"type": item_type,
"asset_id": asset_id,
"frame_id": frame_id,
@@ -143,14 +162,14 @@ def _insert_item(
"width": width,
"height": height,
"rotation": 0,
- "z_index": z_index if z_index is not None else _next_z(conn),
+ "z_index": z_index if z_index is not None else _next_z(conn, surface),
"created_at": now,
"updated_at": now,
}
conn.execute(
f"INSERT INTO moodboard_items ({_ITEM_COLUMNS}) VALUES "
- "(:id, :project_id, :type, :asset_id, :frame_id, :parent_id, :data, :x, :y, :width, "
- ":height, :rotation, :z_index, :created_at, :updated_at)",
+ "(:id, :project_id, :surface, :type, :asset_id, :frame_id, :parent_id, :data, :x, :y, "
+ ":width, :height, :rotation, :z_index, :created_at, :updated_at)",
item,
)
return get_item(conn, item["id"])
@@ -209,6 +228,57 @@ def add_gen_node(
)
+#: Trainer-canvas nodes. They only ever exist on the `trainer` surface (the Trainer tab's graph);
+#: `resource` is the exception - a read-only telemetry node usable on either canvas.
+TRAINER_SURFACE = "trainer"
+
+
+def add_train_dataset(
+ conn: sqlite3.Connection, x: float, y: float, surface: str = TRAINER_SURFACE
+) -> dict[str, Any]:
+ return _insert_item(
+ conn, item_type="trainDataset", x=x, y=y, width=280, height=240,
+ data={"datasetId": None}, surface=surface,
+ )
+
+
+def add_caption(
+ conn: sqlite3.Connection, x: float, y: float, surface: str = TRAINER_SURFACE
+) -> dict[str, Any]:
+ return _insert_item(
+ conn, item_type="caption", x=x, y=y, width=280, height=200,
+ data={"datasetId": None, "overwrite": False}, surface=surface,
+ )
+
+
+def add_trainer(
+ conn: sqlite3.Connection, x: float, y: float, surface: str = TRAINER_SURFACE
+) -> dict[str, Any]:
+ return _insert_item(
+ # Roomier than the other nodes on purpose: the body is a live log tail.
+ conn, item_type="trainer", x=x, y=y, width=420, height=340,
+ # `runId` is persisted so the node rebinds to its run after a reload and can still Resume.
+ data={"datasetId": None, "runId": None, "hyperparams": {}}, surface=surface,
+ )
+
+
+def add_loss_graph(
+ conn: sqlite3.Connection, x: float, y: float, surface: str = TRAINER_SURFACE
+) -> dict[str, Any]:
+ return _insert_item(
+ conn, item_type="lossGraph", x=x, y=y, width=320, height=220,
+ data={"runId": None}, surface=surface,
+ )
+
+
+def add_resource(
+ conn: sqlite3.Connection, x: float, y: float, surface: str = STUDIO_SURFACE
+) -> dict[str, Any]:
+ return _insert_item(
+ conn, item_type="resource", x=x, y=y, width=280, height=170, data={}, surface=surface
+ )
+
+
def add_layer(conn: sqlite3.Connection, x: float, y: float) -> dict[str, Any]:
return _insert_item(
conn, item_type="layer", x=x, y=y, width=420, height=300, data={"name": "Layer"}, z_index=0
@@ -306,20 +376,29 @@ def delete_item(conn: sqlite3.Connection, item_id: str) -> None:
def replace_board(
- conn: sqlite3.Connection, items: list[dict[str, Any]], connectors: list[dict[str, Any]]
+ conn: sqlite3.Connection,
+ items: list[dict[str, Any]],
+ connectors: list[dict[str, Any]],
+ surface: str = STUDIO_SURFACE,
) -> None:
- """Replace the whole board (undo/redo restore), preserving ids, in one transaction."""
+ """Replace ONE surface's board (undo/redo restore), preserving ids, in one transaction. Scoped
+ by surface so a Studio undo never wipes the Trainer canvas (they share the table)."""
pid = _project_id(conn)
- conn.execute("DELETE FROM moodboard_connectors WHERE project_id = ?", (pid,))
- conn.execute("DELETE FROM moodboard_items WHERE project_id = ?", (pid,))
+ conn.execute(
+ "DELETE FROM moodboard_connectors WHERE project_id = ? AND surface = ?", (pid, surface)
+ )
+ conn.execute(
+ "DELETE FROM moodboard_items WHERE project_id = ? AND surface = ?", (pid, surface)
+ )
for it in items:
conn.execute(
f"INSERT INTO moodboard_items ({_ITEM_COLUMNS}) VALUES "
- "(:id, :project_id, :type, :asset_id, :frame_id, :parent_id, :data, :x, :y, :width, "
- ":height, :rotation, :z_index, :created_at, :updated_at)",
+ "(:id, :project_id, :surface, :type, :asset_id, :frame_id, :parent_id, :data, :x, :y, "
+ ":width, :height, :rotation, :z_index, :created_at, :updated_at)",
{
"id": it["id"],
"project_id": pid,
+ "surface": surface,
"type": it["type"],
"asset_id": it.get("assetId"),
"frame_id": it.get("frameId"),
@@ -338,11 +417,12 @@ def replace_board(
for c in connectors:
conn.execute(
"INSERT INTO moodboard_connectors "
- "(id, project_id, from_item_id, to_item_id, label, data, created_at) "
- "VALUES (:id, :project_id, :from_item_id, :to_item_id, :label, :data, :created_at)",
+ "(id, project_id, surface, from_item_id, to_item_id, label, data, created_at) VALUES "
+ "(:id, :project_id, :surface, :from_item_id, :to_item_id, :label, :data, :created_at)",
{
"id": c["id"],
"project_id": pid,
+ "surface": surface,
"from_item_id": c["fromItemId"],
"to_item_id": c["toItemId"],
"label": c.get("label"),
@@ -359,11 +439,14 @@ def create_connector(
source_handle: str | None = None,
target_handle: str | None = None,
) -> dict[str, Any]:
- get_item(conn, from_item_id)
+ source = get_item(conn, from_item_id)
get_item(conn, to_item_id)
connector = {
"id": _uuid(),
"projectId": _project_id(conn),
+ # A connector lives on the same canvas as the nodes it joins, so it's derived rather than
+ # passed in - the caller never has to know which surface it's drawing on.
+ "surface": source["surface"],
"fromItemId": from_item_id,
"toItemId": to_item_id,
"label": None,
@@ -372,11 +455,11 @@ def create_connector(
}
conn.execute(
"INSERT INTO moodboard_connectors "
- "(id, project_id, from_item_id, to_item_id, label, data, created_at) "
- "VALUES (?, ?, ?, ?, ?, ?, ?)",
+ "(id, project_id, surface, from_item_id, to_item_id, label, data, created_at) "
+ "VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
(
- connector["id"], connector["projectId"], from_item_id, to_item_id, None,
- json.dumps(connector["data"]), connector["createdAt"],
+ connector["id"], connector["projectId"], connector["surface"], from_item_id,
+ to_item_id, None, json.dumps(connector["data"]), connector["createdAt"],
),
)
return connector
diff --git a/core/src/inline_core/studio/schema.py b/core/src/inline_core/studio/schema.py
index 0ed6705..b5c8b36 100644
--- a/core/src/inline_core/studio/schema.py
+++ b/core/src/inline_core/studio/schema.py
@@ -1,16 +1,18 @@
"""SQLite schema for a project's ``project.db`` - a faithful port of the Studio TypeScript
-``electron/main/db/schema.ts`` (SCHEMA_VERSION 14). The DB is the source of truth for a project;
-"save" is implicit. Bumping ``SCHEMA_VERSION`` + adding a migration is how the schema evolves.
+``electron/main/db/schema.ts`` (ported at SCHEMA_VERSION 14). The DB is the source of truth for a
+project; "save" is implicit. Bump ``SCHEMA_VERSION`` + add a migration to evolve the schema.
Kept byte-compatible with the Node schema so Core can open existing ``.inlinestudio`` projects: same
tables, same column names, same ``user_version`` stamping, and the same additive/rename migrations.
+v15 adds the Core-only LoRA training tables (``training_datasets``/``_dataset_items``/``_runs``);
+they are purely additive, so an older project gains them on next open.
"""
from __future__ import annotations
import sqlite3
-SCHEMA_VERSION = 14
+SCHEMA_VERSION = 16
SCHEMA_SQL = """
CREATE TABLE IF NOT EXISTS project (
@@ -86,6 +88,9 @@
CREATE TABLE IF NOT EXISTS moodboard_items (
id TEXT PRIMARY KEY,
project_id TEXT NOT NULL,
+ -- Which canvas the item lives on: the Studio moodboard or the Trainer tab's graph. One table,
+ -- two surfaces, so both canvases reuse the same item/connector plumbing.
+ surface TEXT NOT NULL DEFAULT 'studio',
type TEXT NOT NULL DEFAULT 'asset',
asset_id TEXT,
frame_id TEXT,
@@ -104,6 +109,7 @@
CREATE TABLE IF NOT EXISTS moodboard_connectors (
id TEXT PRIMARY KEY,
project_id TEXT NOT NULL,
+ surface TEXT NOT NULL DEFAULT 'studio',
from_item_id TEXT NOT NULL,
to_item_id TEXT NOT NULL,
label TEXT,
@@ -137,8 +143,50 @@
CREATE INDEX IF NOT EXISTS idx_assets_project ON assets(project_id);
CREATE INDEX IF NOT EXISTS idx_assets_folder ON assets(folder_id);
CREATE INDEX IF NOT EXISTS idx_asset_folders_parent ON asset_folders(parent_id);
+CREATE TABLE IF NOT EXISTS training_datasets (
+ id TEXT PRIMARY KEY,
+ project_id TEXT NOT NULL,
+ name TEXT NOT NULL,
+ trigger_word TEXT NOT NULL DEFAULT '',
+ created_at INTEGER NOT NULL,
+ updated_at INTEGER NOT NULL
+);
+
+CREATE TABLE IF NOT EXISTS training_dataset_items (
+ id TEXT PRIMARY KEY,
+ dataset_id TEXT NOT NULL,
+ asset_id TEXT NOT NULL,
+ caption TEXT NOT NULL DEFAULT '',
+ position INTEGER NOT NULL,
+ created_at INTEGER NOT NULL
+);
+
+CREATE TABLE IF NOT EXISTS training_runs (
+ id TEXT PRIMARY KEY,
+ project_id TEXT NOT NULL,
+ dataset_id TEXT NOT NULL,
+ name TEXT NOT NULL,
+ status TEXT NOT NULL DEFAULT 'queued',
+ base_mode TEXT NOT NULL DEFAULT 'turbo_adapter',
+ hyperparams TEXT NOT NULL DEFAULT '{}',
+ output_lora_path TEXT,
+ progress_fraction REAL NOT NULL DEFAULT 0,
+ progress_status TEXT,
+ step INTEGER NOT NULL DEFAULT 0,
+ total_steps INTEGER NOT NULL DEFAULT 0,
+ checkpoint_path TEXT,
+ gpu_ids TEXT NOT NULL DEFAULT '[]',
+ error TEXT,
+ created_at INTEGER NOT NULL,
+ updated_at INTEGER NOT NULL
+);
+
CREATE INDEX IF NOT EXISTS idx_moodboard_items_project ON moodboard_items(project_id);
CREATE INDEX IF NOT EXISTS idx_moodboard_connectors_project ON moodboard_connectors(project_id);
+CREATE INDEX IF NOT EXISTS idx_training_datasets_project ON training_datasets(project_id);
+CREATE INDEX IF NOT EXISTS idx_training_dataset_items_dataset ON training_dataset_items(dataset_id);
+CREATE INDEX IF NOT EXISTS idx_training_runs_project ON training_runs(project_id);
+CREATE INDEX IF NOT EXISTS idx_training_runs_dataset ON training_runs(dataset_id);
"""
@@ -187,6 +235,14 @@ def _migrate_columns(conn: sqlite3.Connection) -> None:
_add_column_if_missing(conn, "moodboard_items", "frame_id", "TEXT")
_add_column_if_missing(conn, "moodboard_items", "parent_id", "TEXT")
+
+ # v15 -> v16: the Trainer tab gets its own canvas; existing items stay on the Studio surface.
+ _add_column_if_missing(
+ conn, "moodboard_items", "surface", "TEXT NOT NULL DEFAULT 'studio'"
+ )
+ _add_column_if_missing(
+ conn, "moodboard_connectors", "surface", "TEXT NOT NULL DEFAULT 'studio'"
+ )
_add_column_if_missing(conn, "frame_inputs", "source_frame_id", "TEXT")
_relax_frame_inputs_asset_id(conn) # v8 -> v9: asset_id must be nullable
diff --git a/core/src/inline_core/studio/system_stats.py b/core/src/inline_core/studio/system_stats.py
new file mode 100644
index 0000000..a9c4ff2
--- /dev/null
+++ b/core/src/inline_core/studio/system_stats.py
@@ -0,0 +1,92 @@
+"""Periodic host + GPU telemetry, broadcast on ``events:systemStats`` for the Trainer tab.
+
+Samples CPU/RAM (psutil) and per-GPU utilization/VRAM/temp (pynvml) about once a second and fans it
+out to every ``/events`` subscriber. Skips sampling entirely when nobody is listening, so it costs
+nothing with no client connected. GPU stats are NVIDIA-only; an MPS/CPU host reports host metrics
+with an empty ``gpus`` list. Missing psutil/pynvml degrade to zeros rather than erroring.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+_INTERVAL_SECONDS = 1.5
+
+
+class SystemStats:
+ """A background poller that broadcasts ``SystemStatsEvent`` frames while clients connect."""
+
+ def __init__(self, events: Any) -> None:
+ self._events = events
+ self._task: asyncio.Task[None] | None = None
+ self._nvml_ready = False
+
+ def start(self) -> None:
+ if self._task is None:
+ self._task = asyncio.create_task(self._loop())
+
+ def stop(self) -> None:
+ if self._task is not None:
+ self._task.cancel()
+ self._task = None
+
+ async def _loop(self) -> None:
+ try:
+ while True:
+ await asyncio.sleep(_INTERVAL_SECONDS)
+ if self._events.subscriber_count == 0:
+ continue # nobody watching - don't spend a sample
+ stats = await asyncio.to_thread(self._sample)
+ self._events.broadcast("events:systemStats", stats)
+ except asyncio.CancelledError:
+ pass
+
+ def _sample(self) -> dict[str, Any]:
+ cpu, ram_used, ram_total = self._host()
+ return {"cpu": cpu, "ramUsed": ram_used, "ramTotal": ram_total, "gpus": self._gpus()}
+
+ def _host(self) -> tuple[float, int, int]:
+ try:
+ import psutil
+
+ vm = psutil.virtual_memory()
+ return float(psutil.cpu_percent()), int(vm.used), int(vm.total)
+ except Exception: # noqa: BLE001 - psutil optional; telemetry degrades to zeros
+ return 0.0, 0, 0
+
+ def _gpus(self) -> list[dict[str, Any]]:
+ try:
+ import pynvml
+ except Exception: # noqa: BLE001 - no NVML (non-NVIDIA or extra not installed)
+ return []
+ try:
+ if not self._nvml_ready:
+ pynvml.nvmlInit()
+ self._nvml_ready = True
+ out: list[dict[str, Any]] = []
+ for i in range(pynvml.nvmlDeviceGetCount()):
+ handle = pynvml.nvmlDeviceGetHandleByIndex(i)
+ mem = pynvml.nvmlDeviceGetMemoryInfo(handle)
+ util = pynvml.nvmlDeviceGetUtilizationRates(handle)
+ name = pynvml.nvmlDeviceGetName(handle)
+ out.append(
+ {
+ "index": i,
+ "name": name.decode() if isinstance(name, bytes) else name,
+ "utilization": float(util.gpu),
+ "memoryUsed": int(mem.used),
+ "memoryTotal": int(mem.total),
+ "temperature": _temperature(pynvml, handle),
+ }
+ )
+ return out
+ except Exception: # noqa: BLE001 - any NVML hiccup: report no GPUs rather than crash
+ return []
+
+
+def _temperature(pynvml: Any, handle: Any) -> float | None:
+ try:
+ return float(pynvml.nvmlDeviceGetTemperature(handle, pynvml.NVML_TEMPERATURE_GPU))
+ except Exception: # noqa: BLE001
+ return None
diff --git a/core/src/inline_core/studio/training.py b/core/src/inline_core/studio/training.py
new file mode 100644
index 0000000..cee9cab
--- /dev/null
+++ b/core/src/inline_core/studio/training.py
@@ -0,0 +1,458 @@
+"""Orchestrate LoRA training: a Studio-side long-running job that spawns the trainer as a subprocess
+and streams progress over ``/events`` - mirroring the ffmpeg timeline (``timeline/render.py``).
+
+Training never runs inside Core's graph executor (``core/CLAUDE.md``: the executor never runs the
+denoise loop). ``start`` writes the dataset to a working dir + a manifest, launches
+``python -m inline_core.training `` (or ``accelerate launch --multi_gpu`` for 2+ GPUs),
+and returns immediately; the subprocess emits JSON-line progress that this class parses, mirrors
+into the durable ``training_runs`` row, and broadcasts. The produced ``.safetensors`` is written
+into ``models_dir()/loras/`` so it shows in the LoRA loader node's dropdown - no loader changes.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import re
+import shutil
+import sys
+import uuid
+from pathlib import Path
+from typing import Any
+
+from ..config import models_dir
+from . import training_store as ts
+
+_CAPTION_EXT = ".txt"
+
+
+def _safe(name: str) -> str:
+ return re.sub(r"[^\w.-]+", "_", name).strip("_") or "lora"
+
+
+def _unique_lora_name(loras: Path, stem: str) -> str:
+ """`loras/.safetensors`, suffixed -2, -3… if that file already exists. Retraining under
+ the same name shouldn't silently overwrite a LoRA the user may already be generating with."""
+ if not (loras / f"{stem}.safetensors").exists():
+ return f"loras/{stem}.safetensors"
+ for n in range(2, 1000):
+ if not (loras / f"{stem}-{n}.safetensors").exists():
+ return f"loras/{stem}-{n}.safetensors"
+ return f"loras/{stem}-{uuid.uuid4().hex[:6]}.safetensors"
+
+
+class Training:
+ """Owns the LoRA training lifecycle: dataset CRUD, auto-caption, and one run at a time."""
+
+ def __init__(self, store: Any, events: Any, on_output: Any = None) -> None:
+ self._store = store
+ self._events = events
+ # Rescans the model catalog when a run's LoRA lands in models/loras/, so it appears in the
+ # loader node's dropdown + bumps the registry version - the same hook model downloads use.
+ self._on_output = on_output
+ # One run holds the GPU at a time; the process is kept so cancel can SIGTERM it.
+ self._active: dict[str, asyncio.subprocess.Process] = {}
+ self._cancelled: set[str] = set()
+
+ def _conn(self) -> Any:
+ return self._store.conn()
+
+ # --- dataset CRUD (delegates to training_store) ---------------------------------------------
+
+ def list_datasets(self) -> list[dict[str, Any]]:
+ return ts.list_datasets(self._conn())
+
+ def create_dataset(self, inp: dict[str, Any]) -> dict[str, Any]:
+ return ts.create_dataset(self._conn(), inp["name"], inp.get("triggerWord") or "")
+
+ def list_items(self, dataset_id: str) -> list[dict[str, Any]]:
+ return ts.list_items(self._conn(), dataset_id)
+
+ def add_items(self, dataset_id: str, asset_ids: list[str]) -> list[dict[str, Any]]:
+ return ts.add_items(self._conn(), dataset_id, asset_ids)
+
+ def remove_item(self, item_id: str) -> None:
+ ts.remove_item(self._conn(), item_id)
+
+ def set_caption(self, item_id: str, caption: str) -> dict[str, Any]:
+ return ts.set_caption(self._conn(), item_id, caption)
+
+ def list_runs(self) -> list[dict[str, Any]]:
+ conn = self._conn()
+ self._reconcile_orphans(conn)
+ return ts.list_runs(conn)
+
+ def _reconcile_orphans(self, conn: Any) -> None:
+ """Flip runs left mid-flight by a crash/restart to `interrupted`.
+
+ A run row stays `training` if its process died without the orchestrator seeing the exit (a
+ killed subprocess, a server restart). Nothing would ever clear it, and since the UI blocks
+ starting while any run is training, one orphan disables training forever. Checked on the
+ Trainer tab's own load, which is the first thing the UI does."""
+ for run in ts.list_runs(conn):
+ if run["status"] in ("training", "queued") and run["id"] not in self._active:
+ ts.update_run(
+ conn, run["id"],
+ {"status": "interrupted",
+ "error": "Training stopped unexpectedly; you can resume it."},
+ )
+
+ def status(self, run_id: str) -> dict[str, Any]:
+ return ts.get_run(self._conn(), run_id)
+
+ def _persist(self, run_id: str, patch: dict[str, Any]) -> bool:
+ """Write a run patch, tolerating the row having gone away. False when it has.
+
+ `_conn()` resolves the *currently open* project, but a run belongs to the project it was
+ started in - switching projects mid-run makes its row unreachable. That used to raise
+ inside the training task ("Training run not found", an unretrieved task exception) and
+ left the run stuck as `training` forever."""
+ try:
+ ts.update_run(self._conn(), run_id, patch)
+ return True
+ except ValueError:
+ return False
+
+ def _lookup(self, run_id: str) -> dict[str, Any] | None:
+ try:
+ return ts.get_run(self._conn(), run_id)
+ except ValueError:
+ return None
+
+ # --- auto-caption ---------------------------------------------------------------------------
+
+ async def auto_caption(self, dataset_id: str, overwrite: bool) -> list[dict[str, Any]]:
+ """Caption items with the local VLM (a subprocess, so torch never imports server-side)."""
+ conn, folder = self._conn(), self._store.folder()
+ items = ts.list_items(conn, dataset_id)
+ targets = [
+ {"id": it["id"], "path": str(folder / self._asset_path(it["assetId"]))}
+ for it in items
+ if overwrite or not it["caption"].strip()
+ ]
+ if not targets:
+ return items
+ total = len(targets)
+ proc = await asyncio.create_subprocess_exec(
+ sys.executable, "-m", "inline_core.training.caption",
+ stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE,
+ stderr=asyncio.subprocess.STDOUT,
+ )
+ # Feed the manifest and close stdin, then read stdout LINE BY LINE. `communicate()` would
+ # buffer the whole run, so the captioner's per-image lines (which it already emits) only
+ # landed after it finished - no progress. The manifest is small, so writing before reading
+ # can't deadlock.
+ assert proc.stdin is not None and proc.stdout is not None
+ proc.stdin.write(json.dumps({"items": targets}).encode())
+ await proc.stdin.drain()
+ proc.stdin.close()
+
+ self._progress_caption(dataset_id, 0, total)
+ done = 0
+ async for raw in proc.stdout:
+ msg = _parse_json_line(raw.decode(errors="replace"))
+ if msg and msg.get("id") and isinstance(msg.get("caption"), str):
+ ts.set_caption(conn, msg["id"], msg["caption"])
+ done += 1
+ self._progress_caption(dataset_id, done, total, item_id=msg["id"])
+ await proc.wait()
+ self._progress_caption(dataset_id, done, total, finished=True)
+ return ts.list_items(conn, dataset_id)
+
+ def _progress_caption(
+ self,
+ dataset_id: str,
+ done: int,
+ total: int,
+ *,
+ item_id: str | None = None,
+ finished: bool = False,
+ ) -> None:
+ self._events.broadcast(
+ "events:captionProgress",
+ {
+ "datasetId": dataset_id,
+ "done": done,
+ "total": total,
+ "itemId": item_id,
+ "finished": finished,
+ },
+ )
+
+ def _relativize(self, path: str) -> str:
+ try:
+ return str(Path(path).resolve().relative_to(self._store.folder().resolve()))
+ except (ValueError, RuntimeError):
+ return path
+
+ def _asset_path(self, asset_id: str) -> str:
+ from . import assets as ax
+
+ file = ax.asset_file(self._conn(), asset_id)
+ if file is None:
+ raise ValueError("Dataset references a missing asset.")
+ return file["filePath"]
+
+ # --- training run ---------------------------------------------------------------------------
+
+ def start(self, dataset_id: str, hyperparams: dict[str, Any]) -> dict[str, Any]:
+ if self._active:
+ raise RuntimeError("A training run is already in progress; wait for it to finish.")
+ dataset = ts.get_dataset(self._conn(), dataset_id)
+ run = ts.create_run(self._conn(), dataset_id, dataset["name"], hyperparams)
+ asyncio.create_task(self._run(run["id"], resume=False))
+ return run
+
+ def resume(self, run_id: str) -> dict[str, Any]:
+ if self._active:
+ raise RuntimeError("A training run is already in progress; wait for it to finish.")
+ run = ts.get_run(self._conn(), run_id)
+ if run["status"] not in ("interrupted", "failed"):
+ raise RuntimeError("Only an interrupted run can be resumed.")
+ asyncio.create_task(self._run(run_id, resume=True))
+ return ts.update_run(self._conn(), run_id, {"status": "queued", "error": None})
+
+ def cancel(self, run_id: str) -> None:
+ proc = self._active.get(run_id)
+ if proc is not None:
+ self._cancelled.add(run_id)
+ proc.terminate() # SIGTERM -> the trainer flushes a final checkpoint before exit
+
+ async def _run(self, run_id: str, *, resume: bool) -> None:
+ try:
+ manifest_path, output_rel = self._prepare(run_id, resume=resume)
+ except Exception as error: # noqa: BLE001 - surface prep failures as a run error
+ self._persist(run_id, {"status": "failed", "error": str(error)})
+ self._events.broadcast("events:trainingError", {"runId": run_id, "error": str(error)})
+ return
+
+ self._persist(run_id, {"status": "training", "progressStatus": "starting"})
+ proc = await self._spawn(run_id, manifest_path)
+ self._active[run_id] = proc
+ saw_done = False
+ try:
+ saw_done = await self._drain(run_id, proc, output_rel)
+ finally:
+ await proc.wait()
+ self._active.pop(run_id, None)
+ self._finish(run_id, proc.returncode or 0, saw_done, output_rel)
+
+ async def _drain(self, run_id: str, proc: asyncio.subprocess.Process, output_rel: str) -> bool:
+ """Read the trainer's JSON-line stdout, mirroring progress into the row + events."""
+ saw_done = False
+ last_status = ""
+ assert proc.stdout is not None
+ async for raw in proc.stdout:
+ line = _last_progress_segment(raw.decode(errors="replace"))
+ msg = _parse_json_line(line)
+ if msg is None:
+ # Not protocol JSON - the trainer's own stdout/stderr (loader bars, warnings,
+ # tracebacks). Surfaced as log lines so the Trainer node can show them live.
+ if line:
+ self._events.broadcast("events:trainingLog", {"runId": run_id, "line": line})
+ continue
+ kind = msg.get("type")
+ if kind == "progress":
+ self._on_progress(run_id, msg)
+ # Progress arrives as protocol JSON, so it never reached the log pane - the node
+ # showed loader noise but no training steps. Mirror it as a log line: every step
+ # once a loss exists, and each phase change (loading/caching) once.
+ entry = _progress_log_line(msg, last_status)
+ if entry is not None:
+ last_status = str(msg.get("status") or "")
+ self._events.broadcast(
+ "events:trainingLog", {"runId": run_id, "line": entry}
+ )
+ elif kind == "sample" and msg.get("path"):
+ # /media serves project-relative paths, so relativize the trainer's absolute one.
+ self._events.broadcast(
+ "events:trainingSample",
+ {
+ "runId": run_id,
+ "step": int(msg.get("step", 0)),
+ "path": self._relativize(msg["path"]),
+ },
+ )
+ elif kind == "checkpoint" and msg.get("path"):
+ self._persist(run_id, {"checkpointPath": msg["path"]})
+ elif kind == "error" and msg.get("message"):
+ # A trainer-side failure: keep it in the run's log so the node shows why it died.
+ self._events.broadcast(
+ "events:trainingLog", {"runId": run_id, "line": f"error: {msg['message']}"}
+ )
+ elif kind == "done":
+ saw_done = True
+ return saw_done
+
+ def _on_progress(self, run_id: str, msg: dict[str, Any]) -> None:
+ step = int(msg.get("step", 0))
+ total = int(msg.get("total", 0)) or None
+ fraction = float(msg.get("fraction", 0.0))
+ status = msg.get("status")
+ patch: dict[str, Any] = {"progressFraction": fraction, "step": step}
+ if total:
+ patch["totalSteps"] = total
+ if status:
+ patch["progressStatus"] = status
+ self._persist(run_id, patch)
+ event: dict[str, Any] = {
+ "runId": run_id, "fraction": fraction, "step": step, "totalSteps": total or 0,
+ }
+ if isinstance(msg.get("loss"), int | float):
+ event["loss"] = float(msg["loss"])
+ if status:
+ event["status"] = status
+ self._events.broadcast("events:trainingProgress", event)
+
+ def _finish(self, run_id: str, code: int, saw_done: bool, output_rel: str) -> None:
+ if run_id in self._cancelled:
+ self._cancelled.discard(run_id)
+ # SIGTERM flushes a final checkpoint (trainer.py: "a resumable cancel"), so a cancel
+ # that got far enough to checkpoint is offered for Resume (the status the UI shows a
+ # Resume on), not a dead-end "cancelled". A cancel before any checkpoint stays terminal.
+ run = self._lookup(run_id)
+ if run is None:
+ return # its project was closed/switched mid-run; nothing left to record
+ resumable = bool(run.get("checkpointPath")) or run["step"] > 0
+ if resumable:
+ self._persist(
+ run_id,
+ {"status": "interrupted", "progressStatus": "cancelled",
+ "error": "Training was cancelled; you can resume it."},
+ )
+ self._events.broadcast(
+ "events:trainingError", {"runId": run_id, "error": "Cancelled; resumable."}
+ )
+ else:
+ self._persist(run_id, {"status": "cancelled", "progressStatus": "cancelled"})
+ self._events.broadcast(
+ "events:trainingError", {"runId": run_id, "error": "Cancelled."}
+ )
+ return
+ if saw_done and code == 0 and (models_dir() / output_rel).is_file():
+ self._persist(
+ run_id, {"status": "done", "progressFraction": 1.0, "outputLoraPath": output_rel}
+ )
+ if self._on_output is not None:
+ self._on_output() # rescan so the new LoRA shows in the loader dropdown
+ self._events.broadcast(
+ "events:trainingDone", {"runId": run_id, "outputLoraPath": output_rel}
+ )
+ return
+ # Non-zero exit with a checkpoint is resumable; otherwise it failed outright.
+ run = self._lookup(run_id)
+ if run is None:
+ return
+ resumable = bool(run.get("checkpointPath")) or run["step"] > 0
+ status = "interrupted" if resumable else "failed"
+ error = "Training was interrupted; you can resume it." if resumable else "Training failed."
+ self._persist(run_id, {"status": status, "error": error})
+ self._events.broadcast("events:trainingError", {"runId": run_id, "error": error})
+
+ # --- subprocess plumbing --------------------------------------------------------------------
+
+ def _prepare(self, run_id: str, *, resume: bool) -> tuple[Path, str]:
+ """Export the dataset + write the run manifest. Returns (manifest path, loras-rel path)."""
+ conn, folder = self._conn(), self._store.folder()
+ run = ts.get_run(conn, run_id)
+ dataset = ts.get_dataset(conn, run["datasetId"])
+ items = ts.list_items(conn, run["datasetId"])
+ if not items:
+ raise ValueError("Add at least one image to the dataset before training.")
+
+ work = folder / "training_runs" / run_id
+ dataset_dir = work / "dataset"
+ checkpoint_dir = work / "checkpoints"
+ for d in (dataset_dir, checkpoint_dir):
+ d.mkdir(parents=True, exist_ok=True)
+
+ trigger = dataset["triggerWord"].strip()
+ for i, item in enumerate(items):
+ src = folder / self._asset_path(item["assetId"])
+ dest = dataset_dir / f"{i:04d}{src.suffix.lower()}"
+ shutil.copyfile(src, dest)
+ caption = ", ".join(p for p in (trigger, item["caption"].strip()) if p)
+ dest.with_suffix(_CAPTION_EXT).write_text(caption, encoding="utf-8")
+
+ loras = models_dir() / "loras"
+ loras.mkdir(parents=True, exist_ok=True)
+ # A user-chosen name wins; otherwise fall back to "-". The run id keeps
+ # the fallback unique, but a chosen name is used verbatim so the file is findable in the
+ # loader dropdown - de-duplicated only if it would clobber an existing LoRA.
+ chosen = _safe(str(run["hyperparams"].get("outputName") or "").strip())
+ if chosen and chosen != "lora":
+ output_rel = _unique_lora_name(loras, chosen)
+ else:
+ output_rel = f"loras/{_safe(run['name'])}-{run_id[:8]}.safetensors"
+ resume_from = str(checkpoint_dir) if resume else None
+ manifest = {
+ "runId": run_id,
+ "workingDir": str(work),
+ "datasetDir": str(dataset_dir),
+ "checkpointDir": str(checkpoint_dir),
+ "outputPath": str(models_dir() / output_rel),
+ "resumeFrom": resume_from,
+ "modelsDir": str(models_dir()),
+ "baseMode": run["hyperparams"]["baseMode"],
+ "triggerWord": trigger,
+ "hyperparams": run["hyperparams"],
+ "gpuIds": run["hyperparams"].get("gpuIds") or [],
+ }
+ manifest_path = work / "manifest.json"
+ manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8")
+ return manifest_path, output_rel
+
+ async def _spawn(
+ self, run_id: str, manifest_path: Path
+ ) -> asyncio.subprocess.Process:
+ import os
+
+ run = ts.get_run(self._conn(), run_id)
+ gpu_ids = [int(g) for g in (run["hyperparams"].get("gpuIds") or [])]
+ env = {**os.environ, "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"}
+ if gpu_ids:
+ env["CUDA_VISIBLE_DEVICES"] = ",".join(str(g) for g in gpu_ids)
+ base = [sys.executable]
+ if len(gpu_ids) > 1:
+ # DDP: only the small LoRA gradients all-reduce, so scaling is near-linear.
+ base += ["-m", "accelerate.commands.launch", "--multi_gpu",
+ "--num_processes", str(len(gpu_ids))]
+ cmd = [*base, "-m", "inline_core.training", str(manifest_path)]
+ return await asyncio.create_subprocess_exec(
+ *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT, env=env,
+ )
+
+
+def _last_progress_segment(raw: str) -> str:
+ """One clean log line out of a raw stdout chunk.
+
+ Progress bars (tqdm, the weight loaders) redraw with ``\\r`` and no newline, so a whole bar's
+ worth of redraws arrives as a single chunk. Keeping only the last segment yields the bar's
+ current state instead of every intermediate frame concatenated into one enormous line."""
+ return raw.replace("\r\n", "\n").split("\r")[-1].rstrip()
+
+
+def _progress_log_line(msg: dict[str, Any], last_status: str) -> str | None:
+ """A log line for a progress tick, or None when there's nothing new to say.
+
+ Once the loop is running every step carries a loss, so each one is worth a line. Before that
+ the trainer only reports phases (loading encoders, caching latents), which repeat - those are
+ logged once, when the phase changes."""
+ status = str(msg.get("status") or "")
+ loss = msg.get("loss")
+ if isinstance(loss, (int, float)):
+ total = int(msg.get("total", 0))
+ step = int(msg.get("step", 0))
+ return f"step {step}/{total} · loss {float(loss):.4f}"
+ return status if status and status != last_status else None
+
+
+def _parse_json_line(line: str) -> dict[str, Any] | None:
+ line = line.strip()
+ if not line or not line.startswith("{"):
+ return None
+ try:
+ parsed = json.loads(line)
+ except json.JSONDecodeError:
+ return None
+ return parsed if isinstance(parsed, dict) else None
diff --git a/core/src/inline_core/studio/training_store.py b/core/src/inline_core/studio/training_store.py
new file mode 100644
index 0000000..46b9407
--- /dev/null
+++ b/core/src/inline_core/studio/training_store.py
@@ -0,0 +1,268 @@
+"""CRUD for the LoRA-training tables (``training_datasets`` / ``_dataset_items`` / ``_runs``).
+
+Pure DB functions over an open ``sqlite3.Connection``, mirroring ``assets.py``. Row mappers emit the
+camelCase shapes the wire contract (``TrainingDataset`` / ``TrainingDatasetItem`` / ``TrainingRun``
+in ``src/shared/types.ts``) expects. The ``Training`` orchestrator (``training.py``) drives the
+actual run; this module only owns the rows.
+"""
+
+from __future__ import annotations
+
+import json
+import sqlite3
+import time
+import uuid
+from typing import Any
+
+# The hyperparam fields that also get their own columns for querying/display.
+_DEFAULT_HYPERPARAMS: dict[str, Any] = {
+ "baseMode": "turbo_adapter",
+ "rank": 16,
+ "alpha": 16,
+ "learningRate": 1e-4,
+ "steps": 1500,
+ "batchSize": 1,
+ "resolution": 1024,
+ "saveEvery": 250,
+ "gpuIds": [],
+}
+
+
+def _now() -> int:
+ return int(time.time() * 1000)
+
+
+def _uuid() -> str:
+ return str(uuid.uuid4())
+
+
+def _project_id(conn: sqlite3.Connection) -> str:
+ row = conn.execute("SELECT id FROM project LIMIT 1").fetchone()
+ if row is None:
+ raise RuntimeError("No project is open.")
+ return row["id"]
+
+
+# --- datasets -----------------------------------------------------------------------------------
+
+
+def _row_to_dataset(row: sqlite3.Row) -> dict[str, Any]:
+ return {
+ "id": row["id"],
+ "projectId": row["project_id"],
+ "name": row["name"],
+ "triggerWord": row["trigger_word"],
+ "createdAt": row["created_at"],
+ "updatedAt": row["updated_at"],
+ }
+
+
+def list_datasets(conn: sqlite3.Connection) -> list[dict[str, Any]]:
+ rows = conn.execute("SELECT * FROM training_datasets ORDER BY created_at DESC").fetchall()
+ return [_row_to_dataset(r) for r in rows]
+
+
+def get_dataset(conn: sqlite3.Connection, dataset_id: str) -> dict[str, Any]:
+ row = conn.execute("SELECT * FROM training_datasets WHERE id = ?", (dataset_id,)).fetchone()
+ if row is None:
+ raise ValueError("Training dataset not found.")
+ return _row_to_dataset(row)
+
+
+def create_dataset(conn: sqlite3.Connection, name: str, trigger_word: str) -> dict[str, Any]:
+ trimmed = name.strip()
+ if not trimmed:
+ raise ValueError("Dataset name is required.")
+ now = _now()
+ row = {
+ "id": _uuid(),
+ "project_id": _project_id(conn),
+ "name": trimmed,
+ "trigger_word": trigger_word.strip(),
+ "created_at": now,
+ "updated_at": now,
+ }
+ conn.execute(
+ "INSERT INTO training_datasets "
+ "(id, project_id, name, trigger_word, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?)",
+ tuple(row.values()),
+ )
+ return get_dataset(conn, row["id"])
+
+
+def _touch_dataset(conn: sqlite3.Connection, dataset_id: str) -> None:
+ conn.execute(
+ "UPDATE training_datasets SET updated_at = ? WHERE id = ?", (_now(), dataset_id)
+ )
+
+
+# --- dataset items ------------------------------------------------------------------------------
+
+
+def _row_to_item(row: sqlite3.Row) -> dict[str, Any]:
+ return {
+ "id": row["id"],
+ "datasetId": row["dataset_id"],
+ "assetId": row["asset_id"],
+ "caption": row["caption"],
+ "position": row["position"],
+ "createdAt": row["created_at"],
+ }
+
+
+def list_items(conn: sqlite3.Connection, dataset_id: str) -> list[dict[str, Any]]:
+ rows = conn.execute(
+ "SELECT * FROM training_dataset_items WHERE dataset_id = ? ORDER BY position",
+ (dataset_id,),
+ ).fetchall()
+ return [_row_to_item(r) for r in rows]
+
+
+def add_items(
+ conn: sqlite3.Connection, dataset_id: str, asset_ids: list[str]
+) -> list[dict[str, Any]]:
+ """Append library assets as items, skipping any already in the dataset. Returns the new rows."""
+ get_dataset(conn, dataset_id) # validate
+ existing = {
+ r["asset_id"]
+ for r in conn.execute(
+ "SELECT asset_id FROM training_dataset_items WHERE dataset_id = ?", (dataset_id,)
+ ).fetchall()
+ }
+ start = conn.execute(
+ "SELECT COALESCE(MAX(position) + 1, 0) AS n FROM training_dataset_items "
+ "WHERE dataset_id = ?",
+ (dataset_id,),
+ ).fetchone()["n"]
+ created: list[dict[str, Any]] = []
+ position = start
+ for asset_id in asset_ids:
+ if asset_id in existing:
+ continue
+ existing.add(asset_id)
+ item_id = _uuid()
+ conn.execute(
+ "INSERT INTO training_dataset_items (id, dataset_id, asset_id, caption, position, "
+ "created_at) VALUES (?, ?, ?, ?, ?, ?)",
+ (item_id, dataset_id, asset_id, "", position, _now()),
+ )
+ created.append(get_item(conn, item_id))
+ position += 1
+ if created:
+ _touch_dataset(conn, dataset_id)
+ return created
+
+
+def get_item(conn: sqlite3.Connection, item_id: str) -> dict[str, Any]:
+ row = conn.execute(
+ "SELECT * FROM training_dataset_items WHERE id = ?", (item_id,)
+ ).fetchone()
+ if row is None:
+ raise ValueError("Dataset item not found.")
+ return _row_to_item(row)
+
+
+def set_caption(conn: sqlite3.Connection, item_id: str, caption: str) -> dict[str, Any]:
+ item = get_item(conn, item_id)
+ conn.execute(
+ "UPDATE training_dataset_items SET caption = ? WHERE id = ?", (caption, item_id)
+ )
+ _touch_dataset(conn, item["datasetId"])
+ return get_item(conn, item_id)
+
+
+def remove_item(conn: sqlite3.Connection, item_id: str) -> None:
+ item = get_item(conn, item_id)
+ conn.execute("DELETE FROM training_dataset_items WHERE id = ?", (item_id,))
+ _touch_dataset(conn, item["datasetId"])
+
+
+# --- runs ---------------------------------------------------------------------------------------
+
+
+def _row_to_run(row: sqlite3.Row) -> dict[str, Any]:
+ hyperparams = json.loads(row["hyperparams"] or "{}")
+ return {
+ "id": row["id"],
+ "projectId": row["project_id"],
+ "datasetId": row["dataset_id"],
+ "name": row["name"],
+ "status": row["status"],
+ "hyperparams": {**_DEFAULT_HYPERPARAMS, **hyperparams},
+ "outputLoraPath": row["output_lora_path"],
+ "progressFraction": row["progress_fraction"],
+ "progressStatus": row["progress_status"],
+ "step": row["step"],
+ "totalSteps": row["total_steps"],
+ "error": row["error"],
+ "createdAt": row["created_at"],
+ "updatedAt": row["updated_at"],
+ }
+
+
+def list_runs(conn: sqlite3.Connection) -> list[dict[str, Any]]:
+ rows = conn.execute("SELECT * FROM training_runs ORDER BY created_at DESC").fetchall()
+ return [_row_to_run(r) for r in rows]
+
+
+def get_run(conn: sqlite3.Connection, run_id: str) -> dict[str, Any]:
+ row = conn.execute("SELECT * FROM training_runs WHERE id = ?", (run_id,)).fetchone()
+ if row is None:
+ raise ValueError("Training run not found.")
+ return _row_to_run(row)
+
+
+def create_run(
+ conn: sqlite3.Connection, dataset_id: str, name: str, hyperparams: dict[str, Any]
+) -> dict[str, Any]:
+ dataset = get_dataset(conn, dataset_id)
+ merged = {**_DEFAULT_HYPERPARAMS, **hyperparams}
+ now = _now()
+ run_id = _uuid()
+ conn.execute(
+ "INSERT INTO training_runs (id, project_id, dataset_id, name, status, base_mode, "
+ "hyperparams, total_steps, gpu_ids, created_at, updated_at) "
+ "VALUES (?, ?, ?, ?, 'queued', ?, ?, ?, ?, ?, ?)",
+ (
+ run_id,
+ dataset["projectId"],
+ dataset_id,
+ name.strip() or dataset["name"],
+ str(merged["baseMode"]),
+ json.dumps(merged),
+ int(merged["steps"]),
+ json.dumps(merged["gpuIds"]),
+ now,
+ now,
+ ),
+ )
+ return get_run(conn, run_id)
+
+
+# Columns the orchestrator may patch as a run progresses (camelCase key -> column).
+_RUN_PATCH_COLUMNS = {
+ "status": "status",
+ "outputLoraPath": "output_lora_path",
+ "progressFraction": "progress_fraction",
+ "progressStatus": "progress_status",
+ "step": "step",
+ "totalSteps": "total_steps",
+ "checkpointPath": "checkpoint_path",
+ "error": "error",
+}
+
+
+def update_run(conn: sqlite3.Connection, run_id: str, patch: dict[str, Any]) -> dict[str, Any]:
+ sets: list[str] = []
+ values: list[Any] = []
+ for key, value in patch.items():
+ column = _RUN_PATCH_COLUMNS.get(key)
+ if column is None:
+ continue
+ sets.append(f"{column} = ?")
+ values.append(value)
+ sets.append("updated_at = ?")
+ values.append(_now())
+ values.append(run_id)
+ conn.execute(f"UPDATE training_runs SET {', '.join(sets)} WHERE id = ?", tuple(values))
+ return get_run(conn, run_id)
diff --git a/core/src/inline_core/training/__init__.py b/core/src/inline_core/training/__init__.py
new file mode 100644
index 0000000..205e208
--- /dev/null
+++ b/core/src/inline_core/training/__init__.py
@@ -0,0 +1,10 @@
+"""LoRA training subpackage - runs as a **subprocess**, never imported by the server.
+
+The Studio orchestrator (``studio/training.py``) launches ``python -m inline_core.training
+``; this package trains a LoRA and reports progress as JSON lines (``protocol.py``). Heavy
+deps (torch/diffusers/peft) live behind the ``training`` optional-deps extra and import lazily in
+``trainer.py`` / ``dataset.py`` / ``caption.py`` - importing this package is cheap + torch-free, so
+the server never pulls the training stack in.
+"""
+
+from __future__ import annotations
diff --git a/core/src/inline_core/training/__main__.py b/core/src/inline_core/training/__main__.py
new file mode 100644
index 0000000..81577e4
--- /dev/null
+++ b/core/src/inline_core/training/__main__.py
@@ -0,0 +1,67 @@
+"""Trainer entry point: ``python -m inline_core.training ``.
+
+Reads the run manifest the orchestrator wrote, runs the LoRA training loop, and reports progress as
+JSON lines on stdout (``protocol.py``). The heavy ``trainer`` import is deferred to ``main`` so a
+bad invocation - or an install without the ``training`` extra - reports a clean error instead of an
+``ImportError`` traceback. A SIGTERM (the orchestrator's cancel) asks the loop to flush a final
+checkpoint and stop; the run is then resumable.
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+from pathlib import Path
+from typing import Any
+
+from . import protocol
+
+
+def _message_for(error: Exception, manifest: dict[str, Any]) -> str:
+ """Turn a CUDA OOM into something the user can act on. Training resolution drives peak VRAM far
+ more than rank does, so point at that first - a 16GB card trains Z-Image at 512 but not 768."""
+ if type(error).__name__ != "OutOfMemoryError":
+ return str(error)
+ resolution = int(manifest.get("hyperparams", {}).get("resolution", 0) or 0)
+ lower = " Try 512." if resolution > 512 else ""
+ at = f" at {resolution}px" if resolution else ""
+ return (
+ f"Ran out of GPU memory training{at}. Lower the training resolution in the node's Adjust "
+ f"panel (it drives peak VRAM far more than rank).{lower}"
+ )
+
+
+def main(argv: list[str]) -> int:
+ if not argv:
+ protocol.error("No manifest path given.")
+ return 2
+ try:
+ manifest = json.loads(Path(argv[0]).read_text(encoding="utf-8"))
+ except (OSError, json.JSONDecodeError) as exc:
+ protocol.error(f"Could not read manifest: {exc}")
+ return 2
+
+ try:
+ from .trainer import train # heavy deps (torch/diffusers/peft) load here, not at import
+ except Exception as exc: # noqa: BLE001 - report a missing training stack cleanly
+ protocol.error(f"Training stack unavailable: {exc}. Install the 'training' extra.")
+ return 1
+
+ try:
+ output = train(manifest)
+ except KeyboardInterrupt:
+ return 130
+ except Exception as exc: # noqa: BLE001 - surface any training failure as one error line
+ protocol.error(_message_for(exc, manifest))
+ return 1
+
+ if output is None:
+ # A cooperative stop (cancel/SIGTERM): a checkpoint was saved, no final LoRA. Non-`done`
+ # exit tells the orchestrator this run is cancelled/resumable, not complete.
+ return 0
+ protocol.done(output)
+ return 0
+
+
+if __name__ == "__main__":
+ raise SystemExit(main(sys.argv[1:]))
diff --git a/core/src/inline_core/training/caption.py b/core/src/inline_core/training/caption.py
new file mode 100644
index 0000000..d5d7c5f
--- /dev/null
+++ b/core/src/inline_core/training/caption.py
@@ -0,0 +1,127 @@
+"""Local auto-captioner: ``python -m inline_core.training.caption`` (reads a manifest on stdin).
+
+The orchestrator (``studio/training.py`` ``auto_caption``) pipes ``{"items": [{"id", "path"}]}`` in
+and reads ``{"id", "caption"}`` JSON lines back. Runs as a subprocess so transformers/torch never
+import server-side. Uses BLIP by default; ``INLINE_CAPTIONER_MODEL`` overrides it. The weights are
+fetched once into the HF cache (opt-in, only when the user hits "Auto-caption"), the same
+fetch-once posture the Z-Image runner uses for its reference components.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+from typing import Any
+
+from . import protocol
+
+#: BLIP ships *inside* transformers (no `trust_remote_code`), so it survives runtime upgrades.
+#: Florence-2 was the original default but its remote code breaks on current transformers
+#: (`Florence2LanguageConfig has no attribute forced_bos_token_id`); it still works through
+#: ``INLINE_CAPTIONER_MODEL`` if you pin a compatible transformers - see `_caption_one`.
+_DEFAULT_MODEL = "Salesforce/blip-image-captioning-large"
+_TASK = ""
+
+
+def _load(model_id: str, token: Any = None) -> tuple[Any, Any, Any]:
+ import torch
+ import transformers
+ from transformers import AutoModelForCausalLM, AutoProcessor
+
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+ dtype = torch.float16 if device == "cuda" else torch.float32
+ kwargs: dict[str, Any] = {"torch_dtype": dtype, "trust_remote_code": True, "token": token}
+ # The in-transformers captioners (BLIP et al): `AutoModelForVision2Seq` was renamed
+ # `AutoModelForImageTextToText` in transformers 5, so resolve whichever this build has.
+ image_text = getattr(transformers, "AutoModelForImageTextToText", None) or getattr(
+ transformers, "AutoModelForVision2Seq", None
+ )
+ try:
+ if image_text is None:
+ raise ValueError("no image-text auto class")
+ model = image_text.from_pretrained(model_id, **kwargs)
+ except (ValueError, KeyError):
+ # Florence-2 and friends register themselves as causal LMs via their remote code.
+ model = AutoModelForCausalLM.from_pretrained(model_id, **kwargs)
+ model = model.to(device)
+ model.eval()
+ processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True, token=token)
+ return model, processor, device
+
+
+def _looks_like_auth_failure(error: Exception) -> bool:
+ text = str(error).lower()
+ return "401" in text or "not a valid model identifier" in text or "repository not found" in text
+
+
+def _load_with_fallback(model_id: str) -> tuple[Any, Any, Any]:
+ """Load the captioner, retrying anonymously if a stored HF token is rejected.
+
+ A stale/expired token in the HF cache makes the hub return 401 for *public* repos too, which
+ surfaces as a confusing "not a valid model identifier". The captioner is public, so an
+ anonymous retry recovers instead of failing the whole run."""
+ try:
+ return _load(model_id)
+ except Exception as first: # noqa: BLE001 - retried anonymously below, then re-raised
+ if not _looks_like_auth_failure(first):
+ raise
+ try:
+ return _load(model_id, token=False)
+ except Exception: # noqa: BLE001 - report the original failure, it's the informative one
+ raise first from None
+
+
+def _caption_one(model: Any, processor: Any, device: str, path: str) -> str:
+ """One caption. Handles both shapes: task-token models (Florence-2, which post-processes a
+ tagged string) and plain image-captioning models (BLIP), which just decode the output."""
+ import torch
+ from PIL import Image
+
+ image = Image.open(path).convert("RGB")
+ task_style = hasattr(processor, "post_process_generation")
+ inputs = (
+ processor(text=_TASK, images=image, return_tensors="pt")
+ if task_style
+ else processor(images=image, return_tensors="pt")
+ ).to(device, model.dtype)
+ with torch.no_grad():
+ ids = model.generate(**inputs, max_new_tokens=128, num_beams=3, do_sample=False)
+ if not task_style:
+ return str(processor.decode(ids[0], skip_special_tokens=True)).strip()
+ text = processor.batch_decode(ids, skip_special_tokens=False)[0]
+ parsed = processor.post_process_generation(
+ text, task=_TASK, image_size=(image.width, image.height)
+ )
+ return str(parsed.get(_TASK, "")).strip()
+
+
+def main() -> int:
+ try:
+ manifest = json.loads(sys.stdin.read() or "{}")
+ except json.JSONDecodeError as exc:
+ protocol.error(f"Bad caption manifest: {exc}")
+ return 2
+ items = manifest.get("items") or []
+ if not items:
+ return 0
+
+ model_id = os.environ.get("INLINE_CAPTIONER_MODEL") or _DEFAULT_MODEL
+ try:
+ model, processor, device = _load_with_fallback(model_id)
+ except Exception as exc: # noqa: BLE001 - a missing captioner degrades to no captions
+ protocol.error(f"Captioner unavailable: {exc}")
+ return 1
+
+ for item in items:
+ try:
+ caption = _caption_one(model, processor, device, item["path"])
+ except Exception as exc: # noqa: BLE001 - one bad image shouldn't sink the batch
+ protocol.error(f"Caption failed for {item.get('id')}: {exc}")
+ continue
+ protocol.emit({"id": item["id"], "caption": caption})
+ return 0
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/core/src/inline_core/training/dataset.py b/core/src/inline_core/training/dataset.py
new file mode 100644
index 0000000..784f361
--- /dev/null
+++ b/core/src/inline_core/training/dataset.py
@@ -0,0 +1,100 @@
+"""Dataset precache: encode every training image to a VAE latent and every caption to text
+embeddings **once**, up front, so the VAE + text encoder can be freed and the loop only touches the
+transformer + LoRA (the big low-VRAM win the plan calls for).
+
+The orchestrator exports the dataset as ``NNNN.`` + ``NNNN.txt`` (caption) pairs. For the MVP
+we resize to a square ``resolution`` (center-crop); aspect-ratio bucketing (ai-toolkit
+``buckets.py``) is a follow-up. The conditioning tensor mirrors ``ZImagePipeline._encode_prompt``
+exactly (chat template + penultimate hidden layer + padding stripped) - see ``_encode_caption``.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+_IMAGE_SUFFIXES = (".png", ".jpg", ".jpeg", ".webp", ".bmp")
+
+
+def _pairs(dataset_dir: Path) -> list[tuple[Path, str]]:
+ out: list[tuple[Path, str]] = []
+ for img in sorted(dataset_dir.iterdir()):
+ if img.suffix.lower() not in _IMAGE_SUFFIXES:
+ continue
+ caption_file = img.with_suffix(".txt")
+ caption = caption_file.read_text(encoding="utf-8").strip() if caption_file.exists() else ""
+ out.append((img, caption))
+ return out
+
+
+def _to_tensor(image: Any, resolution: int) -> Any:
+ import torch
+ from PIL import Image
+
+ img = image.convert("RGB")
+ # Center-crop to square, then resize to the training resolution; pixels normalized to [-1, 1].
+ short = min(img.size)
+ left = (img.width - short) // 2
+ top = (img.height - short) // 2
+ img = img.crop((left, top, left + short, top + short)).resize(
+ (resolution, resolution), Image.LANCZOS
+ )
+ import numpy as np
+
+ arr = np.asarray(img, dtype="float32") / 127.5 - 1.0
+ return torch.from_numpy(arr).permute(2, 0, 1) # CHW
+
+
+def precache(
+ dataset_dir: str,
+ components: Any,
+ device: str,
+ dtype: Any,
+ resolution: int,
+) -> list[dict[str, Any]]:
+ """Return ``[{"latent", "embed"}]`` (CPU tensors) for every image/caption pair."""
+ import torch
+ from PIL import Image
+
+ pairs = _pairs(Path(dataset_dir))
+ if not pairs:
+ raise RuntimeError("The exported dataset is empty.")
+
+ vae, text_encoder, tokenizer = (
+ components.vae,
+ components.text_encoder,
+ components.tokenizer,
+ )
+ scaling = float(getattr(vae.config, "scaling_factor", 1.0) or 1.0)
+ shift = float(getattr(vae.config, "shift_factor", 0.0) or 0.0)
+ cached: list[dict[str, Any]] = []
+ for img_path, caption in pairs:
+ pixels = _to_tensor(Image.open(img_path), resolution).to(device, dtype).unsqueeze(0)
+ with torch.no_grad():
+ # Model-space latent: decode is (z / scaling) + shift, so the forward normalization is
+ # (z - shift) * scaling. Dropping the shift skews every training latent off-manifold.
+ latent = (vae.encode(pixels).latent_dist.sample() - shift) * scaling
+ embed = _encode_caption(text_encoder, tokenizer, caption, device)
+ cached.append({"latent": latent.squeeze(0).cpu(), "embed": embed.cpu()})
+ return cached
+
+
+def _encode_caption(text_encoder: Any, tokenizer: Any, caption: str, device: str) -> Any:
+ """Caption -> conditioning, matching ``ZImagePipeline._encode_prompt`` exactly: wrap in the Qwen
+ chat template (thinking on), take the PENULTIMATE hidden layer (not the last), and strip padding
+ to the real tokens (variable-length (seq, dim)). Training conditioning must mirror inference or
+ the LoRA learns against an embedding the generation path never produces."""
+ text = tokenizer.apply_chat_template(
+ [{"role": "user", "content": caption or ""}],
+ tokenize=False,
+ add_generation_prompt=True,
+ enable_thinking=True,
+ )
+ inputs = tokenizer(
+ [text], padding="max_length", max_length=512, truncation=True, return_tensors="pt"
+ ).to(device)
+ mask = inputs.attention_mask.bool()
+ hidden = text_encoder(
+ input_ids=inputs.input_ids, attention_mask=mask, output_hidden_states=True
+ ).hidden_states[-2]
+ return hidden[0][mask[0]] # (real_seq_len, dim), padding removed
diff --git a/core/src/inline_core/training/models.py b/core/src/inline_core/training/models.py
new file mode 100644
index 0000000..5b77297
--- /dev/null
+++ b/core/src/inline_core/training/models.py
@@ -0,0 +1,125 @@
+"""Resolve + load the Z-Image components for training, reusing the inference loaders.
+
+Bring-your-own-weights, same as generation: the base transformer / VAE / text encoder are the single
+files the user already dropped under ``models/`` for Z-Image. Training additionally needs, in Turbo
+mode, the **training/assistant adapter** that undoes turbo distillation during training - we apply
+it by fusing it into the base with the existing LoRA fuser (``models/lora.py``), so the base behaves
+de-distilled while the trainable LoRA learns on top.
+
+NOTE (needs GPU + weights to finalize): the assistant-adapter semantics and the LoRA target modules
+for ``ZImageTransformer2DModel`` should be validated against the ai-toolkit reference
+(``toolkit/assistant_lora.py``, ``extensions_built_in/diffusion_models/z_image/*``).
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+_ARCH = "z-image"
+_WEIGHT_SUFFIXES = (".safetensors", ".ckpt", ".pt", ".pth", ".sft")
+
+
+@dataclass
+class Encoders:
+ """The pieces the one-off latent/caption precache needs. Loaded, used, then freed *before* the
+ transformer loads - on a 15GB card the two together (8GB encoder + 12.3GB transformer) OOM."""
+
+ vae: Any
+ text_encoder: Any
+ tokenizer: Any
+ scheduler: Any
+
+
+def _first_weight(root: Path) -> str | None:
+ if not root.is_dir():
+ return None
+ files = sorted(
+ p for p in root.iterdir() if p.is_file() and p.suffix.lower() in _WEIGHT_SUFFIXES
+ )
+ return str(files[0]) if files else None
+
+
+def _require(root: Path, category: str, env: str) -> str:
+ picked = os.environ.get(env) or _first_weight(root / category)
+ if not picked:
+ raise RuntimeError(
+ f"No {category} weight found under {root / category}. Add the Z-Image {category} file "
+ f"there (or set {env})."
+ )
+ return picked
+
+
+def _adapter_path(root: Path, base_mode: str) -> str | None:
+ if base_mode != "turbo_adapter":
+ return None
+ picked = os.environ.get("INLINE_ZIMAGE_TRAIN_ADAPTER")
+ if not picked:
+ loras = root / "loras"
+ if loras.is_dir():
+ picked = next(
+ (str(p) for p in sorted(loras.iterdir()) if "adapter" in p.name.lower()), None
+ )
+ if not picked:
+ raise RuntimeError(
+ "Turbo mode needs a training adapter to avoid turbo drift. Add the Z-Image training "
+ "adapter to models/loras/ (or set INLINE_ZIMAGE_TRAIN_ADAPTER), or use de-turbo mode."
+ )
+ return picked
+
+
+def compute_dtype() -> Any:
+ import torch
+
+ if torch.cuda.is_available():
+ return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+ return torch.float32
+
+
+def load_encoders(models_dir: str, device: str, dtype: Any) -> Encoders:
+ """The VAE + text encoder + tokenizer + flow-match scheduler, for the precache pass."""
+ from ..models import loaders
+
+ root = Path(models_dir)
+ vae_file = _require(root, "vae", "INLINE_ZIMAGE_VAE")
+ encoder_file = _require(root, "text_encoders", "INLINE_ZIMAGE_TEXT_ENCODER")
+
+ vae = loaders.load_vae(_ARCH, vae_file, dtype, device=device)
+ text_encoder, tokenizer = loaders.load_text_encoder(_ARCH, encoder_file, dtype, device=device)
+ return Encoders(vae, text_encoder, tokenizer, loaders.load_scheduler(_ARCH))
+
+
+def free_encoders(encoders: Encoders) -> None:
+ """Return the VAE + text-encoder VRAM once latents/captions are cached.
+
+ They are dead weight for the training loop, and the transformer alone already fills a 15GB card
+ (12.3GB) - holding both is the OOM. Dropped, not moved to CPU: a 15GB-RAM host can't take the
+ 8GB encoder either. The scheduler is config-only, so it stays."""
+ import gc
+
+ import torch
+
+ from ..models import loaders
+
+ encoders.vae = None
+ encoders.text_encoder = None
+ encoders.tokenizer = None
+ loaders.unload_components(keep_files=set()) # the transformer isn't loaded yet - drop it all
+ gc.collect()
+ if torch.cuda.is_available():
+ torch.cuda.empty_cache()
+
+
+def load_transformer(models_dir: str, base_mode: str, device: str, dtype: Any) -> Any:
+ """The base transformer (grad-enabled, unquantized). In Turbo mode the training adapter is
+ fused in first, so the base behaves de-distilled while the trainable LoRA learns on top."""
+ from ..graph.loader_runners import LoraRef
+ from ..models import loaders
+
+ root = Path(models_dir)
+ diffusion = _require(root, "diffusion_models", "INLINE_ZIMAGE_MODEL")
+ adapter = _adapter_path(root, base_mode)
+ loras: tuple[LoraRef, ...] = (LoraRef(file=adapter, strength=1.0),) if adapter else ()
+ return loaders.load_diffusion(_ARCH, diffusion, dtype, device=device, loras=loras)
diff --git a/core/src/inline_core/training/protocol.py b/core/src/inline_core/training/protocol.py
new file mode 100644
index 0000000..c0f33c7
--- /dev/null
+++ b/core/src/inline_core/training/protocol.py
@@ -0,0 +1,49 @@
+"""The JSON-line progress protocol between the trainer subprocess and the Studio orchestrator.
+
+Every message is one JSON object on its own stdout line. Keep these shapes in sync with the
+orchestrator's parser (``studio/training.py`` ``_drain`` / ``_on_progress``). Torch-free on purpose
+so the entry point can report a clean error even when the training stack is missing.
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+from typing import Any
+
+
+def emit(message: dict[str, Any]) -> None:
+ sys.stdout.write(json.dumps(message) + "\n")
+ sys.stdout.flush()
+
+
+def progress(
+ step: int, total: int, *, loss: float | None = None, status: str | None = None
+) -> None:
+ message: dict[str, Any] = {
+ "type": "progress",
+ "step": step,
+ "total": total,
+ "fraction": (step / total) if total else 0.0,
+ }
+ if loss is not None:
+ message["loss"] = loss
+ if status is not None:
+ message["status"] = status
+ emit(message)
+
+
+def sample(step: int, path: str) -> None:
+ emit({"type": "sample", "step": step, "path": path})
+
+
+def checkpoint(path: str) -> None:
+ emit({"type": "checkpoint", "path": path})
+
+
+def done(output: str) -> None:
+ emit({"type": "done", "output": output})
+
+
+def error(message: str) -> None:
+ emit({"type": "error", "message": message})
diff --git a/core/src/inline_core/training/trainer.py b/core/src/inline_core/training/trainer.py
new file mode 100644
index 0000000..7edc47c
--- /dev/null
+++ b/core/src/inline_core/training/trainer.py
@@ -0,0 +1,236 @@
+"""The LoRA training loop: PEFT adapter on the Z-Image transformer, rectified-flow loss, checkpoint/
+resume, and a loader-compatible ``.safetensors`` at the end.
+
+Uses ``accelerate`` so the same code runs single-GPU or, under ``accelerate launch --multi_gpu``,
+DDP (only the small LoRA grads all-reduce). Progress + samples + checkpoints are reported as JSON
+lines (``protocol.py``). A SIGTERM flushes a checkpoint and returns ``None`` (a resumable cancel).
+
+The three Z-Image specifics were validated against the real ``ZImageTransformer2DModel`` +
+``ZImagePipeline`` on a GPU (see ``_forward``, ``_TARGET_MODULES``, ``_timesteps`` / the loss):
+ 1. ``forward`` takes per-image lists + a normalized timestep and returns ``.sample`` (a list),
+ 2. the LoRA targets are the attention + w1/w2/w3 feed-forward Linears,
+ 3. the flow-match target is (clean - noise) at timestep (1 - sigma).
+"""
+
+from __future__ import annotations
+
+import json
+import signal
+from pathlib import Path
+from typing import Any
+
+from . import dataset as ds
+from . import models, protocol
+
+# Every ZImageTransformerBlock's attention (q/k/v/out) + SwiGLU feed-forward (w1/w2/w3) Linears -
+# confirmed against ZImageTransformer2DModel.named_modules() (34 blocks, 238 Linears). Z-Image's
+# FeedForward is w1/w2/w3, NOT the diffusers-generic ff.net.*, and `to_out.0` matches only the
+# attention output (not the adaLN Linear that also ends in `.0`).
+_TARGET_MODULES = ["to_q", "to_k", "to_v", "to_out.0", "w1", "w2", "w3"]
+
+
+class _Stop:
+ """A SIGTERM latch: the orchestrator's cancel asks the loop to checkpoint and stop."""
+
+ flagged = False
+
+ def __call__(self, *_a: Any) -> None:
+ self.flagged = True
+
+
+def _optimizer(params: list[Any], lr: float) -> Any:
+ try:
+ import bitsandbytes as bnb
+
+ return bnb.optim.Adam8bit(params, lr=lr) # 8-bit Adam keeps optimizer state small
+ except Exception: # noqa: BLE001 - bitsandbytes is optional; fall back to torch AdamW
+ import torch
+
+ return torch.optim.AdamW(params, lr=lr)
+
+
+def _timesteps(device: Any, shift: float) -> Any:
+ import torch
+
+ # Logit-normal sampling of the noise fraction in (0, 1) - denser near the middle, as flow-match
+ # trainers favor - then Z-Image's static resolution shift (scheduler `shift`,
+ # use_dynamic_shifting=False) so the training noise levels match the schedule inference visits.
+ u = torch.sigmoid(torch.randn((), device=device))
+ return shift * u / (1.0 + (shift - 1.0) * u)
+
+
+def _save_checkpoint(
+ accelerator: Any, transformer: Any, optimizer: Any, ckpt_dir: Path, step: int
+) -> None:
+ """Adapter-only checkpoint: the LoRA weights + optimizer + RNG + step - NOT the frozen ~12GB
+ base (which reloads identically from its file every run). ``accelerator.save_state`` would write
+ the whole base each time; a rank-8 adapter is a few MB, so checkpointing stays cheap and resume
+ is exact. Rank-0 writes; the small files are read by every rank on resume."""
+ import torch
+ from peft import get_peft_model_state_dict
+ from safetensors.torch import save_file
+
+ accelerator.wait_for_everyone()
+ if not accelerator.is_main_process:
+ return
+ model = accelerator.unwrap_model(transformer)
+ adapter = {k: v.detach().to("cpu") for k, v in get_peft_model_state_dict(model).items()}
+ save_file(adapter, str(ckpt_dir / "adapter.safetensors"))
+ torch.save(optimizer.state_dict(), ckpt_dir / "optimizer.pt")
+ rng: dict[str, Any] = {"cpu": torch.get_rng_state()}
+ if torch.cuda.is_available():
+ rng["cuda"] = torch.cuda.get_rng_state()
+ torch.save(rng, ckpt_dir / "rng.pt")
+ (ckpt_dir / "step.json").write_text(json.dumps({"step": step}), encoding="utf-8")
+
+
+def _load_checkpoint(accelerator: Any, transformer: Any, optimizer: Any, resume_from: Path) -> int:
+ """Restore an adapter-only checkpoint (see ``_save_checkpoint``) on every rank and return the
+ step to continue from."""
+ import torch
+ from peft import set_peft_model_state_dict
+ from safetensors.torch import load_file
+
+ model = accelerator.unwrap_model(transformer)
+ set_peft_model_state_dict(model, load_file(str(resume_from / "adapter.safetensors")))
+ optimizer.load_state_dict(torch.load(resume_from / "optimizer.pt", weights_only=False))
+ if (resume_from / "rng.pt").is_file():
+ rng = torch.load(resume_from / "rng.pt", weights_only=False)
+ torch.set_rng_state(rng["cpu"])
+ if torch.cuda.is_available() and "cuda" in rng:
+ torch.cuda.set_rng_state(rng["cuda"])
+ return _resume_step(resume_from)
+
+
+def _resume_step(ckpt_dir: Path) -> int:
+ meta = ckpt_dir / "step.json"
+ if meta.exists():
+ try:
+ return int(json.loads(meta.read_text(encoding="utf-8"))["step"])
+ except (json.JSONDecodeError, KeyError, ValueError):
+ return 0
+ return 0
+
+
+def _save_lora(transformer: Any, output_path: str) -> None:
+ """Write the PEFT adapter as safetensors. Its ``base_model.model...lora_A/lora_B`` keys are read
+ directly by the loader's fuser (``models/lora.py`` strips the ``base_model.model.`` prefix)."""
+ import torch
+ from peft import get_peft_model_state_dict
+ from safetensors.torch import save_file
+
+ state = {
+ k: v.detach().to("cpu", dtype=torch.float32).contiguous()
+ for k, v in get_peft_model_state_dict(transformer).items()
+ }
+ Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+ save_file(state, output_path)
+
+
+def _forward(transformer: Any, noisy: Any, t_norm: Any, embed: Any) -> Any:
+ """One prediction from the real ZImageTransformer2DModel, mirroring ZImagePipeline's call.
+
+ The model takes per-image LISTS - latents as (C, F, H, W) with a temporal axis, captions as
+ (seq, dim) - a NORMALIZED timestep (1=clean, 0=noise; the model multiplies by t_scale=1000
+ itself, so we must not pre-scale), and returns per-image latents in ``.sample`` (a list).
+ Returns this single image's (C, H, W) prediction."""
+ out = transformer(
+ [noisy.unsqueeze(1)], # (C, H, W) -> [(C, 1, H, W)]
+ t_norm.reshape(1), # (1,) per-image timestep
+ [embed], # [(seq, dim)]
+ return_dict=True,
+ )
+ sample = out.sample if hasattr(out, "sample") else out[0]
+ return sample[0].squeeze(1) # (C, 1, H, W) -> (C, H, W)
+
+
+def train(manifest: dict[str, Any]) -> str | None:
+ import torch
+ from accelerate import Accelerator
+ from peft import LoraConfig
+
+ hp = manifest["hyperparams"]
+ steps = int(hp["steps"])
+ save_every = max(1, int(hp.get("saveEvery", 250)))
+ resolution = int(hp.get("resolution", 1024))
+ ckpt_dir = Path(manifest["checkpointDir"])
+ ckpt_dir.mkdir(parents=True, exist_ok=True)
+
+ accelerator = Accelerator(gradient_accumulation_steps=max(1, int(hp.get("batchSize", 1))))
+ device = accelerator.device
+ dtype = models.compute_dtype()
+
+ # Two phases, never overlapping: encoders -> precache -> free, THEN the transformer. Loading all
+ # three at once needs ~20GB and OOMs any 16GB card; apart, peak is just the transformer.
+ protocol.progress(0, steps, status="loading encoders")
+ encoders = models.load_encoders(manifest["modelsDir"], str(device), dtype)
+
+ protocol.progress(0, steps, status="caching latents")
+ data = ds.precache(manifest["datasetDir"], encoders, str(device), dtype, resolution)
+ shift = float(encoders.scheduler.config.get("shift", 1.0) or 1.0)
+ models.free_encoders(encoders)
+
+ protocol.progress(0, steps, status="loading model")
+ transformer = models.load_transformer(
+ manifest["modelsDir"], manifest["baseMode"], str(device), dtype
+ )
+ transformer.requires_grad_(False)
+ transformer.add_adapter(
+ LoraConfig(
+ r=int(hp["rank"]),
+ lora_alpha=int(hp.get("alpha") or hp["rank"]),
+ lora_dropout=0.0,
+ target_modules=_TARGET_MODULES,
+ )
+ )
+ if hasattr(transformer, "enable_gradient_checkpointing"):
+ transformer.enable_gradient_checkpointing()
+ lora_params = [p for p in transformer.parameters() if p.requires_grad]
+ optimizer = _optimizer(lora_params, float(hp["learningRate"]))
+ transformer, optimizer = accelerator.prepare(transformer, optimizer)
+
+ start = 0
+ resume_from = manifest.get("resumeFrom")
+ if resume_from and (Path(resume_from) / "adapter.safetensors").exists():
+ start = _load_checkpoint(accelerator, transformer, optimizer, Path(resume_from))
+
+ stop = _Stop()
+ signal.signal(signal.SIGTERM, stop)
+
+ transformer.train()
+ for step in range(start, steps):
+ if stop.flagged:
+ break
+ item = data[step % len(data)]
+ clean = item["latent"].to(device, dtype) # (C, H, W)
+ embed = item["embed"].to(device, dtype) # (seq, dim)
+ noise = torch.randn_like(clean)
+ sigma = _timesteps(device, shift) # scalar noise fraction in (0, 1)
+ noisy = (1 - sigma) * clean + sigma * noise
+ # Z-Image's transformer is trained to predict (clean - noise): the pipeline NEGATES its
+ # output before handing it to FlowMatchEuler as the (noise - clean) velocity, so the raw
+ # output target is the negation of that. Timestep is 1 - sigma (1=clean, 0=noise).
+ target = clean - noise
+
+ with accelerator.accumulate(transformer):
+ pred = _forward(transformer, noisy, 1.0 - sigma, embed)
+ loss = torch.nn.functional.mse_loss(pred.float(), target.float())
+ accelerator.backward(loss)
+ optimizer.step()
+ optimizer.zero_grad()
+
+ done = step + 1
+ protocol.progress(done, steps, loss=float(loss.detach().item()), status="training")
+ if done % save_every == 0 or done == steps:
+ _save_checkpoint(accelerator, transformer, optimizer, ckpt_dir, done)
+ if accelerator.is_main_process:
+ protocol.checkpoint(str(ckpt_dir))
+
+ if stop.flagged:
+ _save_checkpoint(accelerator, transformer, optimizer, ckpt_dir, step)
+ return None # cooperative cancel: a checkpoint exists, so the run is resumable
+
+ accelerator.wait_for_everyone()
+ if accelerator.is_main_process:
+ _save_lora(accelerator.unwrap_model(transformer), manifest["outputPath"])
+ return manifest["outputPath"]
diff --git a/core/tests/test_studio_training.py b/core/tests/test_studio_training.py
new file mode 100644
index 0000000..fc8af43
--- /dev/null
+++ b/core/tests/test_studio_training.py
@@ -0,0 +1,81 @@
+"""LoRA training: dataset/run CRUD, the orchestrator's stdout parsing, and - the real gate - that a
+PEFT adapter's keys are the ones the existing LoRA fuser reads (so the train->load loop closes).
+"""
+
+from __future__ import annotations
+
+import sqlite3
+
+import pytest
+from inline_core.models import lora
+from inline_core.studio import training_store as ts
+from inline_core.studio.schema import apply_schema
+from inline_core.studio.training import _parse_json_line, _safe
+
+
+@pytest.fixture()
+def conn() -> sqlite3.Connection:
+ c = sqlite3.connect(":memory:")
+ c.row_factory = sqlite3.Row
+ apply_schema(c)
+ c.execute("INSERT INTO project (id, name, created_at, updated_at) VALUES ('p', 'P', 0, 0)")
+ for i in range(3):
+ c.execute(
+ "INSERT INTO assets (id, project_id, name, file_path, kind, created_at) "
+ "VALUES (?, 'p', ?, ?, 'image', 0)",
+ (f"a{i}", f"x{i}.png", f"assets/x{i}.png"),
+ )
+ return c
+
+
+def test_dataset_and_item_crud(conn: sqlite3.Connection) -> None:
+ ds = ts.create_dataset(conn, "Hero", "ohwx")
+ assert ds["triggerWord"] == "ohwx"
+ # add_items dedups and assigns sequential positions.
+ items = ts.add_items(conn, ds["id"], ["a0", "a1", "a2", "a0"])
+ assert [it["position"] for it in items] == [0, 1, 2]
+ updated = ts.set_caption(conn, items[0]["id"], "a portrait")
+ assert updated["caption"] == "a portrait"
+ ts.remove_item(conn, items[1]["id"])
+ assert len(ts.list_items(conn, ds["id"])) == 2
+
+
+def test_run_create_and_update(conn: sqlite3.Connection) -> None:
+ ds = ts.create_dataset(conn, "Hero", "")
+ run = ts.create_run(conn, ds["id"], "Hero LoRA", {"rank": 8, "steps": 1200, "gpuIds": [0, 1]})
+ assert run["status"] == "queued"
+ assert run["hyperparams"]["rank"] == 8
+ assert run["totalSteps"] == 1200
+ assert run["hyperparams"]["gpuIds"] == [0, 1]
+
+ patched = ts.update_run(
+ conn, run["id"], {"status": "training", "progressFraction": 0.5, "step": 600}
+ )
+ assert patched["status"] == "training"
+ assert patched["progressFraction"] == 0.5
+ assert patched["step"] == 600
+
+
+def test_parse_json_line_ignores_noise() -> None:
+ assert _parse_json_line('{"type": "progress", "step": 5}') == {"type": "progress", "step": 5}
+ assert _parse_json_line("loading weights...") is None
+ assert _parse_json_line("") is None
+ assert _parse_json_line("[1,2,3]") is None # not an object
+
+
+def test_safe_filename() -> None:
+ assert _safe("My Hero! v2") == "My_Hero_v2"
+ assert _safe("///") == "lora"
+
+
+def test_peft_adapter_keys_are_fuser_compatible() -> None:
+ """A trained LoRA is saved as a PEFT state dict; its keys must be exactly what the fuser groups
+ + matches, or the LoRA would load with 'unmatched key' and never apply."""
+ stem = "base_model.model.transformer_blocks.0.attn.to_q"
+ state = {f"{stem}.lora_A.weight": object(), f"{stem}.lora_B.weight": object()}
+
+ pairs, _alphas = lora._group(state)
+ assert list(pairs) == [stem] # down/up recognized and paired
+
+ # The fuser strips the PEFT `base_model.model.` prefix, yielding the real module path.
+ assert "transformer_blocks.0.attn.to_q" in lora._candidates(stem)
diff --git a/core/uv.lock b/core/uv.lock
index e88011e..b49d9be 100644
--- a/core/uv.lock
+++ b/core/uv.lock
@@ -72,6 +72,23 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/88/c6/92fcd42f1ba33e1184263f25bfabf3d27c383410470f169e4b8163bf9c17/beautifulsoup4-4.15.0-py3-none-any.whl", hash = "sha256:d6f88de62e1d4e38ecb1077eb9724cd0eff29d2a08ca16a401e9b9e93f117cf9", size = 109924, upload-time = "2026-06-07T16:44:21.566Z" },
]
+[[package]]
+name = "bitsandbytes"
+version = "0.49.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
+ { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
+ { name = "packaging" },
+ { name = "torch", version = "2.6.0+cu124", source = { registry = "https://download.pytorch.org/whl/cu124" }, marker = "sys_platform == 'win32'" },
+ { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'win32'" },
+]
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/29/71/acff7af06c818664aa87ff73e17a52c7788ad746b72aea09d3cb8e424348/bitsandbytes-0.49.2-py3-none-manylinux_2_24_aarch64.whl", hash = "sha256:2fc0830c5f7169be36e60e11f2be067c8f812dfcb829801a8703735842450750", size = 31442815, upload-time = "2026-02-16T21:26:06.783Z" },
+ { url = "https://files.pythonhosted.org/packages/19/57/3443d6f183436fbdaf5000aac332c4d5ddb056665d459244a5608e98ae92/bitsandbytes-0.49.2-py3-none-manylinux_2_24_x86_64.whl", hash = "sha256:54b771f06e1a3c73af5c7f16ccf0fc23a846052813d4b008d10cb6e017dd1c8c", size = 60651714, upload-time = "2026-02-16T21:26:11.579Z" },
+ { url = "https://files.pythonhosted.org/packages/b6/d4/501655842ad6771fb077f576d78cbedb5445d15b1c3c91343ed58ca46f0e/bitsandbytes-0.49.2-py3-none-win_amd64.whl", hash = "sha256:2e0ddd09cd778155388023cbe81f00afbb7c000c214caef3ce83386e7144df7d", size = 55372289, upload-time = "2026-02-16T21:26:16.267Z" },
+]
+
[[package]]
name = "certifi"
version = "2026.6.17"
@@ -181,7 +198,7 @@ name = "cuda-bindings"
version = "13.3.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
- { name = "cuda-pathfinder" },
+ { name = "cuda-pathfinder", marker = "sys_platform != 'win32'" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/51/6b/457ca12dad3ee9bfcc9a545cfd6b64b359ba49de40f776f6e028e678f262/cuda_bindings-13.3.1-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c5879712accf6e14bb01aa5e67440eb84998b8d104b509cc7a6dc0b8f656a474", size = 6053539, upload-time = "2026-05-29T23:11:43.19Z" },
@@ -214,43 +231,43 @@ wheels = [
[package.optional-dependencies]
cublas = [
- { name = "nvidia-cublas", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
- { name = "nvidia-cuda-nvrtc", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+ { name = "nvidia-cublas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+ { name = "nvidia-cuda-nvrtc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
]
cudart = [
- { name = "nvidia-cuda-runtime", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+ { name = "nvidia-cuda-runtime", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
]
cufft = [
- { name = "nvidia-cufft", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
- { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+ { name = "nvidia-cufft", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+ { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
]
cufile = [
- { name = "nvidia-cufile", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+ { name = "nvidia-cufile", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
]
cupti = [
- { name = "nvidia-cuda-cupti", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+ { name = "nvidia-cuda-cupti", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
]
curand = [
- { name = "nvidia-curand", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+ { name = "nvidia-curand", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
]
cusolver = [
- { name = "nvidia-cublas", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
- { name = "nvidia-cusolver", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
- { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
- { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+ { name = "nvidia-cublas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+ { name = "nvidia-cusolver", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+ { name = "nvidia-cusparse", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+ { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
]
cusparse = [
- { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
- { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+ { name = "nvidia-cusparse", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+ { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
]
nvjitlink = [
- { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+ { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
]
nvrtc = [
- { name = "nvidia-cuda-nvrtc", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+ { name = "nvidia-cuda-nvrtc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
]
nvtx = [
- { name = "nvidia-nvtx", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+ { name = "nvidia-nvtx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
]
[[package]]
@@ -510,7 +527,7 @@ wheels = [
[[package]]
name = "inline-core"
-version = "1.2.3"
+version = "1.2.31"
source = { editable = "." }
dependencies = [
{ name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
@@ -561,13 +578,24 @@ server = [
{ name = "imageio-ffmpeg" },
{ name = "uvicorn", extra = ["standard"] },
]
+training = [
+ { name = "bitsandbytes", marker = "sys_platform != 'darwin'" },
+ { name = "einops" },
+ { name = "nvidia-ml-py" },
+ { name = "peft" },
+ { name = "pillow" },
+ { name = "psutil" },
+ { name = "timm" },
+]
[package.metadata]
requires-dist = [
{ name = "accelerate", marker = "extra == 'all'", specifier = ">=0.30" },
{ name = "accelerate", marker = "extra == 'runtime'", specifier = ">=0.30" },
+ { name = "bitsandbytes", marker = "sys_platform != 'darwin' and extra == 'training'", specifier = ">=0.43" },
{ name = "diffusers", marker = "extra == 'all'", specifier = ">=0.36" },
{ name = "diffusers", marker = "extra == 'runtime'", specifier = ">=0.36" },
+ { name = "einops", marker = "extra == 'training'", specifier = ">=0.7" },
{ name = "fastapi", marker = "extra == 'all'", specifier = ">=0.110" },
{ name = "fastapi", marker = "extra == 'server'", specifier = ">=0.110" },
{ name = "httpx", marker = "extra == 'dev'", specifier = ">=0.27" },
@@ -577,7 +605,11 @@ requires-dist = [
{ name = "imageio-ffmpeg", marker = "extra == 'server'", specifier = ">=0.4" },
{ name = "numpy", specifier = ">=1.26" },
{ name = "nvidia-ml-py", marker = "extra == 'parallel'", specifier = ">=12" },
+ { name = "nvidia-ml-py", marker = "extra == 'training'", specifier = ">=12" },
+ { name = "peft", marker = "extra == 'training'", specifier = ">=0.11" },
+ { name = "pillow", marker = "extra == 'training'", specifier = ">=10" },
{ name = "psutil", specifier = ">=5.9" },
+ { name = "psutil", marker = "extra == 'training'", specifier = ">=5.9" },
{ name = "pyright", marker = "extra == 'dev'", specifier = ">=1.1" },
{ name = "pytest", marker = "extra == 'dev'", specifier = ">=8" },
{ name = "ruff", marker = "extra == 'dev'", specifier = ">=0.6" },
@@ -585,6 +617,7 @@ requires-dist = [
{ name = "safetensors", marker = "extra == 'runtime'", specifier = ">=0.4" },
{ name = "scipy", marker = "extra == 'all'", specifier = ">=1.11" },
{ name = "scipy", marker = "extra == 'runtime'", specifier = ">=1.11" },
+ { name = "timm", marker = "extra == 'training'", specifier = ">=1.0" },
{ name = "torch", marker = "sys_platform == 'win32' and extra == 'all'", specifier = ">=2.2", index = "https://download.pytorch.org/whl/cu124" },
{ name = "torch", marker = "sys_platform == 'win32' and extra == 'runtime'", specifier = ">=2.2", index = "https://download.pytorch.org/whl/cu124" },
{ name = "torch", marker = "sys_platform != 'win32' and extra == 'all'", specifier = ">=2.2" },
@@ -597,7 +630,7 @@ requires-dist = [
{ name = "uvicorn", extras = ["standard"], marker = "extra == 'server'", specifier = ">=0.29" },
{ name = "xfuser", marker = "extra == 'parallel'", specifier = ">=0.4" },
]
-provides-extras = ["runtime", "server", "parallel", "dev", "all"]
+provides-extras = ["runtime", "server", "parallel", "training", "dev", "all"]
[[package]]
name = "jinja2"
@@ -876,7 +909,7 @@ name = "nvidia-cublas"
version = "13.1.1.3"
source = { registry = "https://pypi.org/simple" }
dependencies = [
- { name = "nvidia-cuda-nvrtc" },
+ { name = "nvidia-cuda-nvrtc", marker = "sys_platform != 'win32'" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/a7/a1/0bd24ee8c8d03adac032fd2909426a00c88f8c57961b1277ded97f91119f/nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b7a210458267ac818974c53038fbec2e969d5c99f305ab15c72522fa9f001dd5", size = 542848918, upload-time = "2026-04-08T18:46:22.985Z" },
@@ -915,7 +948,7 @@ name = "nvidia-cudnn-cu13"
version = "9.20.0.48"
source = { registry = "https://pypi.org/simple" }
dependencies = [
- { name = "nvidia-cublas" },
+ { name = "nvidia-cublas", marker = "sys_platform != 'win32'" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/56/c5/83384d846b2fd17c44bd499b36c75a45ed4f095fbbb2252294e89cea5c5c/nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:e31454ae00094b0c55319d9d15b6fa2fc50a9e1c0f5c8c80fb75258234e731e1", size = 444574296, upload-time = "2026-03-09T19:28:27.751Z" },
@@ -927,7 +960,7 @@ name = "nvidia-cufft"
version = "12.0.0.61"
source = { registry = "https://pypi.org/simple" }
dependencies = [
- { name = "nvidia-nvjitlink" },
+ { name = "nvidia-nvjitlink", marker = "sys_platform != 'win32'" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" },
@@ -957,9 +990,9 @@ name = "nvidia-cusolver"
version = "12.0.4.66"
source = { registry = "https://pypi.org/simple" }
dependencies = [
- { name = "nvidia-cublas" },
- { name = "nvidia-cusparse" },
- { name = "nvidia-nvjitlink" },
+ { name = "nvidia-cublas", marker = "sys_platform != 'win32'" },
+ { name = "nvidia-cusparse", marker = "sys_platform != 'win32'" },
+ { name = "nvidia-nvjitlink", marker = "sys_platform != 'win32'" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" },
@@ -971,7 +1004,7 @@ name = "nvidia-cusparse"
version = "12.6.3.3"
source = { registry = "https://pypi.org/simple" }
dependencies = [
- { name = "nvidia-nvjitlink" },
+ { name = "nvidia-nvjitlink", marker = "sys_platform != 'win32'" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" },
@@ -1041,6 +1074,29 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/df/b2/87e62e8c3e2f4b32e5fe99e0b86d576da1312593b39f47d8ceef365e95ed/packaging-26.2-py3-none-any.whl", hash = "sha256:5fc45236b9446107ff2415ce77c807cee2862cb6fac22b8a73826d0693b0980e", size = 100195, upload-time = "2026-04-24T20:15:22.081Z" },
]
+[[package]]
+name = "peft"
+version = "0.19.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "accelerate" },
+ { name = "huggingface-hub" },
+ { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
+ { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
+ { name = "packaging" },
+ { name = "psutil" },
+ { name = "pyyaml" },
+ { name = "safetensors" },
+ { name = "torch", version = "2.6.0+cu124", source = { registry = "https://download.pytorch.org/whl/cu124" }, marker = "sys_platform == 'win32'" },
+ { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'win32'" },
+ { name = "tqdm" },
+ { name = "transformers" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/86/cf/037f1e3d5186496c05513a6754639e2dab3038a05f384284d49a9bd06a2d/peft-0.19.1.tar.gz", hash = "sha256:0d97542fe96dcdaa20d3b81c06f26f988618f416a73544ab23c3618ccb674a40", size = 763738, upload-time = "2026-04-16T15:46:45.105Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/e8/b6/f54d676ed93cc2dd2234c3b172ea9c8c3d7d29361e66b1b23dec57a67465/peft-0.19.1-py3-none-any.whl", hash = "sha256:2113f72a81621b5913ef28f9022204c742df111890c5f49d812716a4a301e356", size = 680692, upload-time = "2026-04-16T15:46:42.886Z" },
+]
+
[[package]]
name = "pillow"
version = "12.3.0"
@@ -1572,7 +1628,7 @@ resolution-markers = [
"python_full_version < '3.12' and sys_platform != 'win32'",
]
dependencies = [
- { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" } },
+ { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/7a/97/5a3609c4f8d58b039179648e62dd220f89864f56f7357f5d4f45c29eb2cc/scipy-1.17.1.tar.gz", hash = "sha256:95d8e012d8cb8816c226aef832200b1d45109ed4464303e997c5b13122b297c0", size = 30573822, upload-time = "2026-02-23T00:26:24.851Z" }
wheels = [
@@ -1647,7 +1703,7 @@ resolution-markers = [
"python_full_version >= '3.12' and sys_platform != 'win32'",
]
dependencies = [
- { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" } },
+ { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/a7/25/c2700dfaf6442b4effaa91af24ebce5dc9d31bb4a69706313aae70d72cd0/scipy-1.18.0.tar.gz", hash = "sha256:67b2ad2ad54c72ca6d04975a9b2df8c3638c34ddd5b28738e94fc2b57929d378", size = 30774447, upload-time = "2026-06-19T15:01:43.456Z" }
wheels = [
@@ -1792,7 +1848,7 @@ resolution-markers = [
"python_full_version < '3.12' and sys_platform == 'win32'",
]
dependencies = [
- { name = "mpmath" },
+ { name = "mpmath", marker = "sys_platform == 'win32'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/ca/99/5a5b6f19ff9f083671ddf7b9632028436167cd3d33e11015754e41b249a4/sympy-1.13.1.tar.gz", hash = "sha256:9cebf7e04ff162015ce31c9c6c9144daa34a93bd082f54fd8f12deca4f47515f", size = 7533040, upload-time = "2024-07-19T09:26:51.238Z" }
wheels = [
@@ -1808,13 +1864,31 @@ resolution-markers = [
"python_full_version < '3.12' and sys_platform != 'win32'",
]
dependencies = [
- { name = "mpmath" },
+ { name = "mpmath", marker = "sys_platform != 'win32'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" },
]
+[[package]]
+name = "timm"
+version = "1.0.28"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "huggingface-hub" },
+ { name = "pyyaml" },
+ { name = "safetensors" },
+ { name = "torch", version = "2.6.0+cu124", source = { registry = "https://download.pytorch.org/whl/cu124" }, marker = "sys_platform == 'win32'" },
+ { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'win32'" },
+ { name = "torchvision", version = "0.21.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'win32'" },
+ { name = "torchvision", version = "0.28.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'win32'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/35/03/e41389ac641747bfec48d016fde8be1eade1901e6f2c1aedcb0c8cb4b5d9/timm-1.0.28.tar.gz", hash = "sha256:3789d313fdd5541a327b60180d70dbb4bdec73db8ff0655e413db3c3d134a9a4", size = 2451413, upload-time = "2026-07-11T17:24:32.615Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/c1/76/de1bfac17d183c49c6d0887903d3064ced51cf1d9ba7a8d611c1a8808c4f/timm-1.0.28-py3-none-any.whl", hash = "sha256:e577b88da96b3a722ea5e2f042455ce6f715d398304d8e63b17d126ed7d89968", size = 2597944, upload-time = "2026-07-11T17:24:30.869Z" },
+]
+
[[package]]
name = "tokenizers"
version = "0.22.2"
@@ -1850,13 +1924,13 @@ resolution-markers = [
"python_full_version < '3.12' and sys_platform == 'win32'",
]
dependencies = [
- { name = "filelock" },
- { name = "fsspec" },
- { name = "jinja2" },
- { name = "networkx" },
- { name = "setuptools", marker = "python_full_version >= '3.12'" },
- { name = "sympy", version = "1.13.1", source = { registry = "https://pypi.org/simple" } },
- { name = "typing-extensions" },
+ { name = "filelock", marker = "sys_platform == 'win32'" },
+ { name = "fsspec", marker = "sys_platform == 'win32'" },
+ { name = "jinja2", marker = "sys_platform == 'win32'" },
+ { name = "networkx", marker = "sys_platform == 'win32'" },
+ { name = "setuptools", marker = "python_full_version >= '3.12' and sys_platform == 'win32'" },
+ { name = "sympy", version = "1.13.1", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'win32'" },
+ { name = "typing-extensions", marker = "sys_platform == 'win32'" },
]
wheels = [
{ url = "https://download-r2.pytorch.org/whl/cu124/torch-2.6.0%2Bcu124-cp311-cp311-win_amd64.whl", hash = "sha256:6a1fb2714e9323f11edb6e8abf7aad5f79e45ad25c081cde87681a18d99c29eb", upload-time = "2025-01-30T00:55:31Z" },
@@ -1875,18 +1949,18 @@ resolution-markers = [
dependencies = [
{ name = "cuda-bindings", marker = "python_full_version < '3.15' and sys_platform == 'linux'" },
{ name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "sys_platform == 'linux'" },
- { name = "filelock" },
- { name = "fsspec" },
- { name = "jinja2" },
- { name = "networkx" },
+ { name = "filelock", marker = "sys_platform != 'win32'" },
+ { name = "fsspec", marker = "sys_platform != 'win32'" },
+ { name = "jinja2", marker = "sys_platform != 'win32'" },
+ { name = "networkx", marker = "sys_platform != 'win32'" },
{ name = "nvidia-cudnn-cu13", marker = "sys_platform == 'linux'" },
{ name = "nvidia-cusparselt-cu13", marker = "sys_platform == 'linux'" },
{ name = "nvidia-nccl-cu13", marker = "sys_platform == 'linux'" },
{ name = "nvidia-nvshmem-cu13", marker = "sys_platform == 'linux'" },
- { name = "setuptools" },
- { name = "sympy", version = "1.14.0", source = { registry = "https://pypi.org/simple" } },
+ { name = "setuptools", marker = "sys_platform != 'win32'" },
+ { name = "sympy", version = "1.14.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'win32'" },
{ name = "triton", marker = "python_full_version < '3.15' and sys_platform == 'linux'" },
- { name = "typing-extensions" },
+ { name = "typing-extensions", marker = "sys_platform != 'win32'" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/5b/fe/cba54dc58523434919b66f13a667e36e436deddd77ca519e96553617d4ec/torch-2.13.0-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:e76f9bcecc52b8ff711239a2f7547d5353df95878ab232f0773c1d95928b92f8", size = 111187938, upload-time = "2026-07-08T16:05:17.065Z" },
@@ -1915,6 +1989,58 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/c9/37/ef37ca885265e5f79a168616767dd416a3cea1cc3b28bb6b503ce4a5b652/torchao-0.17.0-py3-none-any.whl", hash = "sha256:02eba449036715b9ae784fbaa1a6f97994bb7b0421ce92d1d5d1c08e5bd6d349", size = 1200680, upload-time = "2026-03-30T22:25:54.457Z" },
]
+[[package]]
+name = "torchvision"
+version = "0.21.0"
+source = { registry = "https://pypi.org/simple" }
+resolution-markers = [
+ "python_full_version >= '3.12' and sys_platform == 'win32'",
+ "python_full_version < '3.12' and sys_platform == 'win32'",
+]
+dependencies = [
+ { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12' and sys_platform == 'win32'" },
+ { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' and sys_platform == 'win32'" },
+ { name = "pillow", marker = "sys_platform == 'win32'" },
+ { name = "torch", version = "2.6.0+cu124", source = { registry = "https://download.pytorch.org/whl/cu124" }, marker = "sys_platform == 'win32'" },
+]
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/88/53/4ad334b9b1d8dd99836869fec139cb74a27781298360b91b9506c53f1d10/torchvision-0.21.0-cp311-cp311-win_amd64.whl", hash = "sha256:49bcfad8cfe2c27dee116c45d4f866d7974bcf14a5a9fbef893635deae322f2f", size = 1560523, upload-time = "2025-01-29T16:28:48.751Z" },
+ { url = "https://files.pythonhosted.org/packages/4c/6a/c7752603060d076dfed95135b78b047dc71792630cbcb022e3693d6f32ef/torchvision-0.21.0-cp312-cp312-win_amd64.whl", hash = "sha256:6eb75d41e3bbfc2f7642d0abba9383cc9ae6c5a4ca8d6b00628c225e1eaa63b3", size = 1560520, upload-time = "2025-01-29T16:28:42.122Z" },
+ { url = "https://files.pythonhosted.org/packages/ed/b4/fc60e3bc003879d3de842baea258fffc3586f4b49cd435a5ba1e09c33315/torchvision-0.21.0-cp313-cp313-win_amd64.whl", hash = "sha256:9147f5e096a9270684e3befdee350f3cacafd48e0c54ab195f45790a9c146d67", size = 1560519, upload-time = "2025-01-29T16:28:22.527Z" },
+]
+
+[[package]]
+name = "torchvision"
+version = "0.28.0"
+source = { registry = "https://pypi.org/simple" }
+resolution-markers = [
+ "python_full_version >= '3.12' and sys_platform != 'win32'",
+ "python_full_version < '3.12' and sys_platform != 'win32'",
+]
+dependencies = [
+ { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12' and sys_platform != 'win32'" },
+ { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' and sys_platform != 'win32'" },
+ { name = "pillow", marker = "sys_platform != 'win32'" },
+ { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'win32'" },
+]
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/7a/b2/1e010052079e4c577007b789db336ea7075f1a426e84d17121fbc3745516/torchvision-0.28.0-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:83fe6c020866a85acd7d97deccc45ff11d66daf42916d04396a4309c66c0ccb8", size = 1856017, upload-time = "2026-07-08T16:07:55.533Z" },
+ { url = "https://files.pythonhosted.org/packages/27/be/1b9c5de9c655ca2df4a74100fa671a7b848532ff787e077ccde14a7dea2a/torchvision-0.28.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:5a38bc6da3d72621be003400b66f66a2b4c6d644fde05f680c2cb7ca8cf8dd6c", size = 7841822, upload-time = "2026-07-08T16:07:49.207Z" },
+ { url = "https://files.pythonhosted.org/packages/0b/9b/f1e68e861d4462e3e195a642c2b448e7b7d3fad5f209487162b9a2133d9b/torchvision-0.28.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:7e80f543b22503d9415e126db5f0ff3917036925e38560ee6b9ae38c571a4002", size = 7670718, upload-time = "2026-07-08T16:07:46.525Z" },
+ { url = "https://files.pythonhosted.org/packages/15/49/c1cab1ecbb3ff1a380a3f99283db1dee61b8afe354f6352c643b65937130/torchvision-0.28.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:e9f54c30cd52e3ef7fd034cc69b7bb7e0964e1c8f8743e018ab92e95b40f9eee", size = 1856020, upload-time = "2026-07-08T16:07:52.182Z" },
+ { url = "https://files.pythonhosted.org/packages/f0/4c/95233776e2def960e5abb7a07931230a545f43717a56a1e1140162033598/torchvision-0.28.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:5cf78ebc401ce64ae19b8c55de866bb836797d559a4de9c25ccbe74cfa642d3a", size = 7842127, upload-time = "2026-07-08T16:07:53.446Z" },
+ { url = "https://files.pythonhosted.org/packages/93/e4/e9b2495d0d57b9f60d63c57d0a910410a81b4b073bf70917bef815291119/torchvision-0.28.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:028a3d481b37d785605620d7cdad897064c5a55bae2aa1f2658766333e291940", size = 7675040, upload-time = "2026-07-08T16:07:58.017Z" },
+ { url = "https://files.pythonhosted.org/packages/20/55/08a726c14c67b37c8aca04b077766909f1c7ed23f76116884fe63b9bd033/torchvision-0.28.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:d483b4aa3f5237569053f749cd1a2b5bb548ca456e40461a5dd087f21149d123", size = 1856021, upload-time = "2026-07-08T16:07:45.386Z" },
+ { url = "https://files.pythonhosted.org/packages/db/8f/40beacd53809194f5259e590d1afaeaa8ad57da15f77c646e6560bcc4616/torchvision-0.28.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:bb6dd6918460ed89cc7644adcc2402991474d6933cf1ce92b390641cb233fddf", size = 7797014, upload-time = "2026-07-08T16:07:43.04Z" },
+ { url = "https://files.pythonhosted.org/packages/32/db/062cdb5a84380a60439775311fff34d89229760d2a50680393dc18699956/torchvision-0.28.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:ad7b3a439265cc3739a4ab5b4c998c0e38ea99c0ee7ca4dea35c5d0b099ec237", size = 7674669, upload-time = "2026-07-08T16:07:38.91Z" },
+ { url = "https://files.pythonhosted.org/packages/c5/b9/da40eca5bbe9596c12ae9899ab7abaf887f5e20f29d08b924b4633714821/torchvision-0.28.0-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:3bd9dba55224a9db4a2d77f6feaa5651770d8c8e86d3d0ddb0fa6bec54c8712b", size = 1856014, upload-time = "2026-07-08T16:07:44.282Z" },
+ { url = "https://files.pythonhosted.org/packages/06/d6/313aafd3df4eaf5f330211bd4e75b7598bddbfee4f55580d3b58536e1b20/torchvision-0.28.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:89f90e29b0966352811b12589f3a3c61943bf2bb9487b9d7bbec10efb1096bb5", size = 7796873, upload-time = "2026-07-08T16:07:30.907Z" },
+ { url = "https://files.pythonhosted.org/packages/b3/41/31f8e959ab8f942600b6357f8999c21d779d5fd3304b0fd204ff4b518239/torchvision-0.28.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:36beb0782976906069ca03d4c9aacaf4b6b838b06ed6c20960ea9c51cce7acdd", size = 7674634, upload-time = "2026-07-08T16:07:29.657Z" },
+ { url = "https://files.pythonhosted.org/packages/6a/80/822a6163da716f8a78141cf6678d74e26a572285d4ea866ef8aa657bb307/torchvision-0.28.0-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:09ce8f56e81f19b9c378ae7bb109f83f6659fd8bc3cd14241a48e4af46e9ed49", size = 1856011, upload-time = "2026-07-08T16:07:33.404Z" },
+ { url = "https://files.pythonhosted.org/packages/7f/d1/cd3f9463b39a790ec8c0c2f6e6c8061edb1562114d04fcdfa786ed889345/torchvision-0.28.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:62c7d110f86a039245b587e4fae60278c649f3bd42ff79cfbc1178eca4e72542", size = 7796742, upload-time = "2026-07-08T16:07:28.339Z" },
+ { url = "https://files.pythonhosted.org/packages/d9/82/3e0a7ad18e99831e2d7f4713d3be717b7159ff5a920862dd5c23c454aa71/torchvision-0.28.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:904cf89af220f8c6b2ed0296bb5065b474ce43b77558e48b2bf9de8b0ba17204", size = 7675526, upload-time = "2026-07-08T16:07:34.572Z" },
+]
+
[[package]]
name = "tqdm"
version = "4.68.4"
diff --git a/core/webui.sh b/core/webui.sh
index bb6e5b1..9fdc7bb 100755
--- a/core/webui.sh
+++ b/core/webui.sh
@@ -54,7 +54,7 @@ Paths
Setup
--install create ./.venv (via uv) and install, then exit
- --extra NAME add an install extra (repeatable): runtime, parallel, server
+ --extra NAME add an install extra (repeatable): runtime, parallel, server, training
-h, --help show this help
Development
diff --git a/package.json b/package.json
index dbaac03..768d023 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
{
"name": "inline-studio",
- "version": "1.2.31",
+ "version": "1.2.41",
"description": "An experimentation layer for visual artists: a node canvas for building generative pipelines, served by the Inline Core engine.",
"author": "Inline Studio",
"license": "MIT",
diff --git a/packages/frontend/pyproject.toml b/packages/frontend/pyproject.toml
index 696fb47..f1f718b 100644
--- a/packages/frontend/pyproject.toml
+++ b/packages/frontend/pyproject.toml
@@ -1,6 +1,6 @@
[project]
name = "inline-studio-frontend"
-version = "1.2.31"
+version = "1.2.41"
description = "Prebuilt Inline Studio web UI (SPA), served by Inline Core. Mirrors comfyui-frontend-package."
requires-python = ">=3.9"
readme = "README.md"
diff --git a/screenshots/lora-trainer.png b/screenshots/lora-trainer.png
new file mode 100644
index 0000000..66b49b1
Binary files /dev/null and b/screenshots/lora-trainer.png differ
diff --git a/src/renderer/components/icons.tsx b/src/renderer/components/icons.tsx
index a178b15..8a0caeb 100644
--- a/src/renderer/components/icons.tsx
+++ b/src/renderer/components/icons.tsx
@@ -189,3 +189,48 @@ export function ChevronDownIcon({ className }: { className?: string }): React.JS
)
}
+
+/** `model_training` - the LoRA Trainer tab. */
+export function TrainIcon({ className }: { className?: string }): React.JSX.Element {
+ return (
+
+ )
+}
+
+/** `dashboard` / grid - a training dataset (a stack of images). */
+export function DatasetIcon({ className }: { className?: string }): React.JSX.Element {
+ return (
+
+ )
+}
+
+/** `inventory_2` - the Trainer sidebar's Outputs tab (produced LoRAs + resumable runs). */
+export function LoraOutputIcon({ className }: { className?: string }): React.JSX.Element {
+ return (
+
+ )
+}
+
+/** `dashboard_customize` - the Studio (node canvas) tab. */
+export function StudioIcon({ className }: { className?: string }): React.JSX.Element {
+ return (
+
+ )
+}
+
+/** `play_arrow` - run / start. */
+export function PlayIcon({ className }: { className?: string }): React.JSX.Element {
+ return (
+
+ )
+}
diff --git a/src/renderer/store/moodboardStore.ts b/src/renderer/store/moodboardStore.ts
index 68e7fc3..9564f09 100644
--- a/src/renderer/store/moodboardStore.ts
+++ b/src/renderer/store/moodboardStore.ts
@@ -54,6 +54,8 @@ interface MoodboardState {
setLoaderHero: (itemId: string, assetId: string) => Promise
/** Add a Preview node. Returns the new item (for connection-drop suggestions). */
addPreview: (x: number, y: number) => Promise
+ /** Utility: a read-only host-telemetry node (no handles). */
+ addResource: (x: number, y: number) => Promise
addLayer: (x: number, y: number) => Promise
addDirector: (x: number, y: number) => Promise
addTrim: (x: number, y: number) => Promise
@@ -351,6 +353,22 @@ export const useMoodboardStore = create((set, get) => ({
}
},
+ addResource: async (x, y) => {
+ try {
+ get().record()
+ const res = await studio().moodboard.addResource(x, y)
+ if (!res.ok) {
+ set({ error: res.error })
+ return null
+ }
+ set((s) => ({ items: [...s.items, res.value] }))
+ return res.value
+ } catch (e) {
+ set({ error: ipcErrorMessage(e) })
+ return null
+ }
+ },
+
addPreview: async (x, y) => {
try {
get().record()
diff --git a/src/renderer/store/trainerBoardStore.ts b/src/renderer/store/trainerBoardStore.ts
new file mode 100644
index 0000000..3a0de98
--- /dev/null
+++ b/src/renderer/store/trainerBoardStore.ts
@@ -0,0 +1,165 @@
+/**
+ * The Trainer tab's canvas state - the same moodboard item/connector model, scoped to the `trainer`
+ * surface so it never mixes with the Studio board.
+ *
+ * Deliberately much smaller than `moodboardStore`: the training graph has no frames, layers, takes
+ * or asset imports, so it only needs load / add / move / connect / delete. Node chrome is shared
+ * (`NodeFrame`, `NodeBadge`) via `BoardActionsProvider`, so the nodes still read as one card family.
+ */
+import { create } from 'zustand'
+import type { MoodboardConnector, MoodboardItem } from '@shared/types'
+import type { MoodboardItemPatch } from '@shared/ipc'
+import { studio } from '@/lib/studio'
+import { ipcErrorMessage } from '../lib/ipcError'
+
+const SURFACE = 'trainer' as const
+
+/** Node kinds the Trainer canvas can add (maps 1:1 to the `moodboard:add*` channels). */
+export type TrainerNodeKind = 'trainDataset' | 'caption' | 'trainer' | 'lossGraph' | 'resource'
+
+interface TrainerBoardState {
+ items: MoodboardItem[]
+ connectors: MoodboardConnector[]
+ loading: boolean
+ error: string | null
+ load: () => Promise
+ addNode: (kind: TrainerNodeKind, x: number, y: number) => Promise
+ updateItem: (id: string, patch: MoodboardItemPatch) => Promise
+ /** Merge into an item's `data` (dataset/run/hyperparam selections live there). */
+ patchData: (id: string, data: Record) => Promise
+ deleteItem: (id: string) => Promise
+ connect: (
+ fromItemId: string,
+ toItemId: string,
+ sourceHandle?: string | null,
+ targetHandle?: string | null,
+ ) => Promise
+ disconnect: (connectorId: string) => Promise
+ setError: (error: string | null) => void
+ /** Node whose settings sidebar is open in the right gutter (params live off the node face). */
+ settingsItemId: string | null
+ toggleSettings: (itemId: string) => void
+}
+
+function addFor(kind: TrainerNodeKind, x: number, y: number) {
+ const m = studio().moodboard
+ switch (kind) {
+ case 'trainDataset':
+ return m.addTrainDataset(x, y)
+ case 'caption':
+ return m.addCaption(x, y)
+ case 'trainer':
+ return m.addTrainer(x, y)
+ case 'lossGraph':
+ return m.addLossGraph(x, y)
+ case 'resource':
+ return m.addResource(x, y, SURFACE)
+ }
+}
+
+export const useTrainerBoardStore = create((set, get) => ({
+ items: [],
+ connectors: [],
+ loading: false,
+ error: null,
+
+ load: async () => {
+ set({ loading: true })
+ try {
+ const res = await studio().moodboard.list(SURFACE)
+ if (!res.ok) return set({ error: res.error, loading: false })
+ set({ items: res.value.items, connectors: res.value.connectors, loading: false })
+ } catch (e) {
+ set({ error: ipcErrorMessage(e), loading: false })
+ }
+ },
+
+ addNode: async (kind, x, y) => {
+ try {
+ const res = await addFor(kind, x, y)
+ if (!res.ok) {
+ set({ error: res.error })
+ return null
+ }
+ set((s) => ({ items: [...s.items, res.value] }))
+ return res.value
+ } catch (e) {
+ set({ error: ipcErrorMessage(e) })
+ return null
+ }
+ },
+
+ updateItem: async (id, patch) => {
+ // Optimistic so dragging stays snappy, then persist.
+ set((s) => ({
+ items: s.items.map((it) =>
+ it.id === id
+ ? {
+ ...it,
+ x: patch.x ?? it.x,
+ y: patch.y ?? it.y,
+ width: patch.width ?? it.width,
+ height: patch.height ?? it.height,
+ data: patch.data ?? it.data,
+ }
+ : it,
+ ),
+ }))
+ try {
+ const res = await studio().moodboard.updateItem(id, patch)
+ if (!res.ok) set({ error: res.error })
+ } catch (e) {
+ set({ error: ipcErrorMessage(e) })
+ }
+ },
+
+ patchData: async (id, data) => {
+ const item = get().items.find((it) => it.id === id)
+ if (!item) return
+ await get().updateItem(id, { data: { ...item.data, ...data } })
+ },
+
+ deleteItem: async (id) => {
+ set((s) => ({
+ items: s.items.filter((it) => it.id !== id),
+ connectors: s.connectors.filter((c) => c.fromItemId !== id && c.toItemId !== id),
+ }))
+ try {
+ const res = await studio().moodboard.deleteItem(id)
+ if (!res.ok) set({ error: res.error })
+ } catch (e) {
+ set({ error: ipcErrorMessage(e) })
+ }
+ },
+
+ connect: async (fromItemId, toItemId, sourceHandle = null, targetHandle = null) => {
+ try {
+ const res = await studio().moodboard.createConnector(
+ fromItemId,
+ toItemId,
+ sourceHandle,
+ targetHandle,
+ )
+ if (!res.ok) return set({ error: res.error })
+ set((s) => ({ connectors: [...s.connectors, res.value] }))
+ } catch (e) {
+ set({ error: ipcErrorMessage(e) })
+ }
+ },
+
+ disconnect: async (connectorId) => {
+ set((s) => ({ connectors: s.connectors.filter((c) => c.id !== connectorId) }))
+ try {
+ const res = await studio().moodboard.deleteConnector(connectorId)
+ if (!res.ok) set({ error: res.error })
+ } catch (e) {
+ set({ error: ipcErrorMessage(e) })
+ }
+ },
+
+ setError: (error) => set({ error }),
+
+ settingsItemId: null,
+ toggleSettings: (itemId) =>
+ set((s) => ({ settingsItemId: s.settingsItemId === itemId ? null : itemId })),
+}))
diff --git a/src/renderer/store/trainingStore.test.ts b/src/renderer/store/trainingStore.test.ts
new file mode 100644
index 0000000..7cba985
--- /dev/null
+++ b/src/renderer/store/trainingStore.test.ts
@@ -0,0 +1,89 @@
+import { beforeEach, describe, expect, it } from 'vitest'
+import type { TrainingHyperparams, TrainingRun } from '@shared/types'
+import { useTrainingStore } from './trainingStore'
+
+const HP: TrainingHyperparams = {
+ baseMode: 'turbo_adapter',
+ rank: 16,
+ alpha: 16,
+ learningRate: 1e-4,
+ steps: 100,
+ batchSize: 1,
+ resolution: 1024,
+ saveEvery: 50,
+ gpuIds: [],
+}
+
+function run(id: string): TrainingRun {
+ return {
+ id,
+ projectId: 'p',
+ datasetId: 'd',
+ name: id,
+ status: 'training',
+ hyperparams: HP,
+ outputLoraPath: null,
+ progressFraction: 0,
+ progressStatus: null,
+ step: 0,
+ totalSteps: 100,
+ error: null,
+ createdAt: 0,
+ updatedAt: 0,
+ }
+}
+
+beforeEach(() => {
+ useTrainingStore.setState({
+ progressByRun: {},
+ lossByRun: {},
+ samplesByRun: {},
+ systemStats: null,
+ runs: [],
+ })
+})
+
+describe('trainingStore reducers', () => {
+ it('records progress and accumulates the loss curve', () => {
+ const s = useTrainingStore.getState()
+ s.applyProgress({ runId: 'r', fraction: 0.1, step: 10, totalSteps: 100, loss: 0.5 })
+ s.applyProgress({ runId: 'r', fraction: 0.2, step: 20, totalSteps: 100, loss: 0.4 })
+ const st = useTrainingStore.getState()
+ expect(st.progressByRun['r']).toEqual({
+ fraction: 0.2,
+ step: 20,
+ totalSteps: 100,
+ status: undefined,
+ })
+ expect(st.lossByRun['r']).toEqual([0.5, 0.4])
+ })
+
+ it('does not push a loss point when the event has none', () => {
+ useTrainingStore
+ .getState()
+ .applyProgress({ runId: 'r', fraction: 0.3, step: 30, totalSteps: 100 })
+ expect(useTrainingStore.getState().lossByRun['r']).toBeUndefined()
+ })
+
+ it('appends sample previews per run', () => {
+ const s = useTrainingStore.getState()
+ s.applySample({ runId: 'r', step: 50, path: 'training_runs/r/samples/1.png' })
+ s.applySample({ runId: 'r', step: 100, path: 'training_runs/r/samples/2.png' })
+ expect(useTrainingStore.getState().samplesByRun['r']).toEqual([
+ 'training_runs/r/samples/1.png',
+ 'training_runs/r/samples/2.png',
+ ])
+ })
+
+ it('stores host/GPU telemetry', () => {
+ const stats = { cpu: 12, ramUsed: 1, ramTotal: 2, gpus: [] }
+ useTrainingStore.getState().applyStats(stats)
+ expect(useTrainingStore.getState().systemStats).toEqual(stats)
+ })
+
+ it('patches a run error in place', () => {
+ useTrainingStore.setState({ runs: [run('r')] })
+ useTrainingStore.getState().applyError({ runId: 'r', error: 'boom' })
+ expect(useTrainingStore.getState().runs[0]!.error).toBe('boom')
+ })
+})
diff --git a/src/renderer/store/trainingStore.ts b/src/renderer/store/trainingStore.ts
new file mode 100644
index 0000000..c38c9d1
--- /dev/null
+++ b/src/renderer/store/trainingStore.ts
@@ -0,0 +1,255 @@
+/**
+ * State for the LoRA Trainer tab: datasets, their items, training runs, and the live run telemetry.
+ * Mutations go through `studio().training.*`; live progress/samples/host-stats arrive via the
+ * training events (subscribed by `subscribeTrainingEvents`, wired in TrainerPanel) which call the
+ * setters here. Mirrors `generationStore`'s shape.
+ */
+import { create } from 'zustand'
+import type {
+ SystemStatsEvent,
+ TrainingDataset,
+ TrainingDatasetItem,
+ TrainingDoneEvent,
+ TrainingErrorEvent,
+ TrainingHyperparams,
+ TrainingProgressEvent,
+ TrainingRun,
+ CaptionProgressEvent,
+ TrainingLogEvent,
+ TrainingSampleEvent,
+} from '@shared/types'
+import { studio } from '@/lib/studio'
+import { ipcErrorMessage } from '../lib/ipcError'
+
+/** Live progress for one run, updated from `onTrainingProgress`. */
+export interface RunProgress {
+ fraction: number
+ step: number
+ totalSteps: number
+ status?: string
+}
+
+interface TrainingState {
+ datasets: TrainingDataset[]
+ itemsByDataset: Record
+ runs: TrainingRun[]
+ activeDatasetId: string | null
+ captioning: boolean
+ error: string | null
+
+ progressByRun: Record
+ lossByRun: Record
+ /** The trainer subprocess's streamed stdout, per run - shown in the Trainer node. */
+ logsByRun: Record
+ /** Live auto-caption progress per dataset, cleared when the captioner exits. */
+ captionProgress: Record
+ samplesByRun: Record
+ systemStats: SystemStatsEvent | null
+
+ loadDatasets: () => Promise
+ createDataset: (name: string, triggerWord: string) => Promise
+ selectDataset: (datasetId: string | null) => void
+ loadItems: (datasetId: string) => Promise
+ addItems: (datasetId: string, assetIds: string[]) => Promise
+ removeItem: (datasetId: string, itemId: string) => Promise
+ setCaption: (datasetId: string, itemId: string, caption: string) => Promise
+ autoCaption: (datasetId: string, overwrite: boolean) => Promise
+ loadRuns: () => Promise
+ /** Returns the created run so a canvas node can persist its `runId` and rebind after a reload. */
+ start: (datasetId: string, hyperparams: TrainingHyperparams) => Promise
+ resume: (runId: string) => Promise
+ cancel: (runId: string) => Promise
+
+ applyProgress: (e: TrainingProgressEvent) => void
+ applySample: (e: TrainingSampleEvent) => void
+ applyLog: (e: TrainingLogEvent) => void
+ applyCaptionProgress: (e: CaptionProgressEvent) => void
+ applyDone: (e: TrainingDoneEvent) => void
+ applyError: (e: TrainingErrorEvent) => void
+ applyStats: (e: SystemStatsEvent) => void
+ setError: (error: string | null) => void
+}
+
+export const useTrainingStore = create((set, get) => ({
+ datasets: [],
+ itemsByDataset: {},
+ runs: [],
+ activeDatasetId: null,
+ captioning: false,
+ error: null,
+ progressByRun: {},
+ lossByRun: {},
+ samplesByRun: {},
+ logsByRun: {},
+ captionProgress: {},
+ systemStats: null,
+
+ loadDatasets: async () => {
+ const res = await studio().training.listDatasets()
+ if (res.ok) set({ datasets: res.value })
+ else set({ error: res.error })
+ },
+
+ createDataset: async (name, triggerWord) => {
+ const res = await studio().training.createDataset({ name, triggerWord })
+ if (!res.ok) {
+ set({ error: res.error })
+ return null
+ }
+ set((s) => ({ datasets: [res.value, ...s.datasets], activeDatasetId: res.value.id }))
+ return res.value
+ },
+
+ selectDataset: (activeDatasetId) => {
+ set({ activeDatasetId })
+ if (activeDatasetId) void get().loadItems(activeDatasetId)
+ },
+
+ loadItems: async (datasetId) => {
+ const res = await studio().training.listItems(datasetId)
+ if (res.ok) set((s) => ({ itemsByDataset: { ...s.itemsByDataset, [datasetId]: res.value } }))
+ else set({ error: res.error })
+ },
+
+ addItems: async (datasetId, assetIds) => {
+ const res = await studio().training.addItems(datasetId, assetIds)
+ if (!res.ok) return set({ error: res.error })
+ await get().loadItems(datasetId)
+ },
+
+ removeItem: async (datasetId, itemId) => {
+ const res = await studio().training.removeItem(itemId)
+ if (!res.ok) return set({ error: res.error })
+ await get().loadItems(datasetId)
+ },
+
+ setCaption: async (datasetId, itemId, caption) => {
+ const res = await studio().training.setCaption(itemId, caption)
+ if (!res.ok) return set({ error: res.error })
+ set((s) => ({
+ itemsByDataset: {
+ ...s.itemsByDataset,
+ [datasetId]: (s.itemsByDataset[datasetId] ?? []).map((it) =>
+ it.id === itemId ? res.value : it,
+ ),
+ },
+ }))
+ },
+
+ autoCaption: async (datasetId, overwrite) => {
+ set({ captioning: true, error: null })
+ try {
+ const res = await studio().training.autoCaption(datasetId, overwrite)
+ if (res.ok) set((s) => ({ itemsByDataset: { ...s.itemsByDataset, [datasetId]: res.value } }))
+ else set({ error: res.error })
+ } catch (e) {
+ set({ error: ipcErrorMessage(e) })
+ } finally {
+ set({ captioning: false })
+ }
+ },
+
+ loadRuns: async () => {
+ const res = await studio().training.listRuns()
+ if (res.ok) set({ runs: res.value })
+ else set({ error: res.error })
+ },
+
+ start: async (datasetId, hyperparams) => {
+ const res = await studio().training.start(datasetId, hyperparams)
+ if (!res.ok) {
+ set({ error: res.error })
+ return null
+ }
+ set((s) => ({ runs: [res.value, ...s.runs.filter((r) => r.id !== res.value.id)] }))
+ return res.value
+ },
+
+ resume: async (runId) => {
+ const res = await studio().training.resume(runId)
+ if (!res.ok) return set({ error: res.error })
+ set((s) => ({ runs: s.runs.map((r) => (r.id === runId ? res.value : r)) }))
+ },
+
+ cancel: async (runId) => {
+ const res = await studio().training.cancel(runId)
+ if (!res.ok) set({ error: res.error })
+ },
+
+ applyProgress: (e) =>
+ set((s) => ({
+ progressByRun: {
+ ...s.progressByRun,
+ [e.runId]: {
+ fraction: e.fraction,
+ step: e.step,
+ totalSteps: e.totalSteps,
+ status: e.status,
+ },
+ },
+ lossByRun:
+ typeof e.loss === 'number'
+ ? { ...s.lossByRun, [e.runId]: [...(s.lossByRun[e.runId] ?? []), e.loss] }
+ : s.lossByRun,
+ })),
+
+ applySample: (e) =>
+ set((s) => ({
+ samplesByRun: { ...s.samplesByRun, [e.runId]: [...(s.samplesByRun[e.runId] ?? []), e.path] },
+ })),
+
+ applyLog: (e) =>
+ set((s) => ({
+ // Capped so a long run can't grow the buffer without bound; the node shows the tail.
+ logsByRun: {
+ ...s.logsByRun,
+ [e.runId]: [...(s.logsByRun[e.runId] ?? []), e.line].slice(-400),
+ },
+ })),
+
+ applyCaptionProgress: (e) =>
+ set((s) => {
+ const next = { ...s.captionProgress }
+ if (e.finished) delete next[e.datasetId]
+ else next[e.datasetId] = { done: e.done, total: e.total }
+ return { captionProgress: next }
+ }),
+
+ applyDone: (e) => {
+ set((s) => ({
+ runs: s.runs.map((r) =>
+ r.id === e.runId
+ ? { ...r, status: 'done', outputLoraPath: e.outputLoraPath, progressFraction: 1 }
+ : r,
+ ),
+ }))
+ void get().loadRuns()
+ },
+
+ applyError: (e) => {
+ set((s) => ({
+ runs: s.runs.map((r) => (r.id === e.runId ? { ...r, error: e.error } : r)),
+ }))
+ // A run ending (cancel/crash) only broadcasts the error - reload so its new status lands,
+ // otherwise a stopped run still reads as "training" and its Stop control never flips back.
+ void get().loadRuns()
+ },
+
+ applyStats: (systemStats) => set({ systemStats }),
+ setError: (error) => set({ error }),
+}))
+
+/** Wire the training + telemetry events to the store. Returns an unsubscribe fn. */
+export function subscribeTrainingEvents(): () => void {
+ const s = useTrainingStore.getState()
+ const unsubs = [
+ studio().events.onTrainingProgress((e) => s.applyProgress(e)),
+ studio().events.onTrainingSample((e) => s.applySample(e)),
+ studio().events.onTrainingLog((e) => s.applyLog(e)),
+ studio().events.onCaptionProgress((e) => s.applyCaptionProgress(e)),
+ studio().events.onTrainingDone((e) => s.applyDone(e)),
+ studio().events.onTrainingError((e) => s.applyError(e)),
+ studio().events.onSystemStats((e) => s.applyStats(e)),
+ ]
+ return () => unsubs.forEach((u) => u())
+}
diff --git a/src/renderer/store/uiStore.ts b/src/renderer/store/uiStore.ts
index d5d9b91..3472c55 100644
--- a/src/renderer/store/uiStore.ts
+++ b/src/renderer/store/uiStore.ts
@@ -1,7 +1,12 @@
/** Workspace-level UI state. */
import { create } from 'zustand'
+/** The two top-level surfaces: the node canvas (Studio) and the LoRA Trainer. */
+export type WorkspaceTab = 'studio' | 'trainer'
+
interface UiState {
+ /** Which top-level tab is showing. */
+ activeTab: WorkspaceTab
/** The frame open in the right-side inspector drawer (null = closed). */
inspectorFrameId: string | null
/** Whether the Settings sidebar (fal API key, etc.) is open. */
@@ -10,6 +15,7 @@ interface UiState {
canvasSelection: string[]
/** Flow-space center of the current canvas viewport (where the user is looking). */
canvasCenter: { x: number; y: number }
+ setActiveTab: (tab: WorkspaceTab) => void
setInspectorFrame: (frameId: string | null) => void
setSettingsOpen: (open: boolean) => void
setCanvasSelection: (ids: string[]) => void
@@ -17,10 +23,12 @@ interface UiState {
}
export const useUiStore = create((set) => ({
+ activeTab: 'studio',
inspectorFrameId: null,
settingsOpen: false,
canvasSelection: [],
canvasCenter: { x: 0, y: 0 },
+ setActiveTab: (activeTab) => set({ activeTab }),
setInspectorFrame: (inspectorFrameId) => set({ inspectorFrameId }),
setSettingsOpen: (settingsOpen) => set({ settingsOpen }),
setCanvasSelection: (canvasSelection) => set({ canvasSelection }),
diff --git a/src/renderer/views/Moodboard/AddNodeMenu.tsx b/src/renderer/views/Moodboard/AddNodeMenu.tsx
index bb83f2b..9b264e6 100644
--- a/src/renderer/views/Moodboard/AddNodeMenu.tsx
+++ b/src/renderer/views/Moodboard/AddNodeMenu.tsx
@@ -18,6 +18,8 @@ interface Entry {
icon: React.JSX.Element
/** Accent the row (the generation node is the AI action). */
accent?: boolean
+ /** Groups the row under a header; ungrouped entries render first, as before. */
+ category?: string
}
const ENTRIES: Entry[] = [
@@ -73,7 +75,7 @@ export function AddNodeMenu({
{/* One scroll area for the whole list (built-ins + fal + Inline Core), not per-section. */}
- 
+ 
