Merge pull request #13 from inlineresearch/z-image-trainer

Z image lora trainer

Author: imprsnst

.nvmrc

@@ -0,0 +1 @@
+22

README.md

@@ -2,12 +2,12 @@
 
 <h3 align="center">AI filmmaking on a node canvas</h3>
 
-<p align="center">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.</p>
+<p align="center">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.</p>
 
 <p align="center">
   <a href="LICENSE"><img alt="License: MIT" src="https://img.shields.io/badge/License-MIT-green?style=for-the-badge"></a>
   <a href="https://www.python.org/downloads/"><img alt="Python 3.11+" src="https://img.shields.io/badge/Python-3.11%2B-blue?style=for-the-badge&logo=python&logoColor=white"></a>
-  <a href="../../releases/latest"><img alt="Latest release" src="https://img.shields.io/badge/Release-v1.2.31-blue?style=for-the-badge"></a>
+  <a href="../../releases/latest"><img alt="Latest release" src="https://img.shields.io/badge/Release-v1.2.41-blue?style=for-the-badge"></a>
   <a href="https://discord.gg/cSUS88VdY9"><img alt="Join our Discord" src="https://img.shields.io/badge/Discord-Join%20the%20community-5865F2?logo=discord&logoColor=white&style=for-the-badge"></a>
 </p>
 
@@ -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.
+
+![Inline Studio Trainer tab showing the LoRA training node graph with a dataset, live logs, and a loss curve](https://raw.githubusercontent.com/inlineresearch/Inline-Studio/main/screenshots/lora-trainer.png)
+
+### 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.
+
+<details>
+<summary><b>Datasets, resuming, trigger words, and base model modes</b></summary>
+
+### 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.
+
+</details>
+
+### 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.

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",

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.

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)

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)