From ca4bdbb14e024cc027ffb2202e71d295633a7ddc Mon Sep 17 00:00:00 2001 From: ashish-aesthisia Date: Wed, 22 Jul 2026 22:34:18 +0530 Subject: [PATCH 1/6] Z-Image LoRA trainer as a second Trainer tab --- .nvmrc | 1 + core/pyproject.toml | 14 + core/src/inline_core/server/app.py | 9 + core/src/inline_core/server/rpc.py | 5 + core/src/inline_core/studio/handlers.py | 20 ++ core/src/inline_core/studio/schema.py | 50 ++- core/src/inline_core/studio/system_stats.py | 92 ++++++ core/src/inline_core/studio/training.py | 302 ++++++++++++++++++ core/src/inline_core/studio/training_store.py | 268 ++++++++++++++++ core/src/inline_core/training/__init__.py | 10 + core/src/inline_core/training/__main__.py | 52 +++ core/src/inline_core/training/caption.py | 86 +++++ core/src/inline_core/training/dataset.py | 80 +++++ core/src/inline_core/training/models.py | 97 ++++++ core/src/inline_core/training/protocol.py | 49 +++ core/src/inline_core/training/trainer.py | 183 +++++++++++ core/tests/test_studio_training.py | 81 +++++ core/uv.lock | 216 ++++++++++--- core/webui.sh | 2 +- src/renderer/components/icons.tsx | 36 +++ src/renderer/store/trainingStore.test.ts | 89 ++++++ src/renderer/store/trainingStore.ts | 217 +++++++++++++ src/renderer/store/uiStore.ts | 8 + .../views/Trainer/DatasetItemsGrid.tsx | 113 +++++++ src/renderer/views/Trainer/HyperparamForm.tsx | 138 ++++++++ src/renderer/views/Trainer/SystemStatsBar.tsx | 62 ++++ src/renderer/views/Trainer/TrainerPanel.tsx | 121 +++++++ .../views/Trainer/TrainingMonitor.tsx | 116 +++++++ src/renderer/views/Workspace/Workspace.tsx | 72 ++++- src/shared/ipc.training.test.ts | 32 ++ src/shared/ipc.ts | 66 ++++ src/shared/types.ts | 124 +++++++ 32 files changed, 2753 insertions(+), 58 deletions(-) create mode 100644 .nvmrc create mode 100644 core/src/inline_core/studio/system_stats.py create mode 100644 core/src/inline_core/studio/training.py create mode 100644 core/src/inline_core/studio/training_store.py create mode 100644 core/src/inline_core/training/__init__.py create mode 100644 core/src/inline_core/training/__main__.py create mode 100644 core/src/inline_core/training/caption.py create mode 100644 core/src/inline_core/training/dataset.py create mode 100644 core/src/inline_core/training/models.py create mode 100644 core/src/inline_core/training/protocol.py create mode 100644 core/src/inline_core/training/trainer.py create mode 100644 core/tests/test_studio_training.py create mode 100644 src/renderer/store/trainingStore.test.ts create mode 100644 src/renderer/store/trainingStore.ts create mode 100644 src/renderer/views/Trainer/DatasetItemsGrid.tsx create mode 100644 src/renderer/views/Trainer/HyperparamForm.tsx create mode 100644 src/renderer/views/Trainer/SystemStatsBar.tsx create mode 100644 src/renderer/views/Trainer/TrainerPanel.tsx create mode 100644 src/renderer/views/Trainer/TrainingMonitor.tsx create mode 100644 src/shared/ipc.training.test.ts diff --git a/.nvmrc b/.nvmrc new file mode 100644 index 0000000..2bd5a0a --- /dev/null +++ b/.nvmrc @@ -0,0 +1 @@ +22 diff --git a/core/pyproject.toml b/core/pyproject.toml index db525b8..c5c3c64 100644 --- a/core/pyproject.toml +++ b/core/pyproject.toml @@ -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..93f4bf2 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), # 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..b7ada54 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: @@ -207,6 +208,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/schema.py b/core/src/inline_core/studio/schema.py index 0ed6705..48705e8 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 = 15 SCHEMA_SQL = """ CREATE TABLE IF NOT EXISTS project ( @@ -137,8 +139,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); """ 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..7023a7c --- /dev/null +++ b/core/src/inline_core/studio/training.py @@ -0,0 +1,302 @@ +"""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 +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" + + +class Training: + """Owns the LoRA training lifecycle: dataset CRUD, auto-caption, and one run at a time.""" + + def __init__(self, store: Any, events: Any) -> None: + self._store = store + self._events = events + # 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]]: + return ts.list_runs(self._conn()) + + def status(self, run_id: str) -> dict[str, Any]: + return ts.get_run(self._conn(), run_id) + + # --- 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 + manifest = {"items": 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, + ) + out, _ = await proc.communicate(json.dumps(manifest).encode()) + for line in out.decode(errors="replace").splitlines(): + msg = _parse_json_line(line) + if msg and msg.get("id") and isinstance(msg.get("caption"), str): + ts.set_caption(conn, msg["id"], msg["caption"]) + return ts.list_items(conn, dataset_id) + + 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: + conn = self._conn() + try: + manifest_path, output_rel = self._prepare(run_id, resume=resume) + except Exception as error: # noqa: BLE001 - surface prep failures as a run error + ts.update_run(conn, run_id, {"status": "failed", "error": str(error)}) + self._events.broadcast("events:trainingError", {"runId": run_id, "error": str(error)}) + return + + ts.update_run(conn, 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 + assert proc.stdout is not None + async for raw in proc.stdout: + msg = _parse_json_line(raw.decode(errors="replace")) + if msg is None: + continue + kind = msg.get("type") + if kind == "progress": + self._on_progress(run_id, msg) + 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"): + ts.update_run(self._conn(), run_id, {"checkpointPath": msg["path"]}) + 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 + ts.update_run(self._conn(), 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: + conn = self._conn() + if run_id in self._cancelled: + self._cancelled.discard(run_id) + ts.update_run(conn, 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(): + ts.update_run( + conn, run_id, + {"status": "done", "progressFraction": 1.0, "outputLoraPath": output_rel}, + ) + 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 = ts.get_run(conn, run_id) + 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." + ts.update_run(conn, 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) + 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 _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..2ce3cc6 --- /dev/null +++ b/core/src/inline_core/training/__main__.py @@ -0,0 +1,52 @@ +"""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 . import protocol + + +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(str(exc)) + 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..ee2c012 --- /dev/null +++ b/core/src/inline_core/training/caption.py @@ -0,0 +1,86 @@ +"""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 Florence-2 (small, ~0.23B) 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 + +_DEFAULT_MODEL = "microsoft/Florence-2-base" +_TASK = "" + + +def _load(model_id: str) -> tuple[Any, Any, Any]: + import torch + from transformers import AutoModelForCausalLM, AutoProcessor + + device = "cuda" if torch.cuda.is_available() else "cpu" + dtype = torch.float16 if device == "cuda" else torch.float32 + model = AutoModelForCausalLM.from_pretrained( + model_id, torch_dtype=dtype, trust_remote_code=True + ).to(device) + model.eval() + processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True) + return model, processor, device + + +def _caption_one(model: Any, processor: Any, device: str, path: str) -> str: + import torch + from PIL import Image + + image = Image.open(path).convert("RGB") + inputs = processor(text=_TASK, images=image, return_tensors="pt").to(device, model.dtype) + with torch.no_grad(): + ids = model.generate( + input_ids=inputs["input_ids"], + pixel_values=inputs["pixel_values"], + max_new_tokens=256, + num_beams=3, + do_sample=False, + ) + 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(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..67d5732 --- /dev/null +++ b/core/src/inline_core/training/dataset.py @@ -0,0 +1,80 @@ +"""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. NOTE (needs GPU + weights): the exact conditioning tensor Z-Image +expects (hidden states vs. pooled) must be validated against the Z-Image runner. +""" + +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) + 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(): + latent = vae.encode(pixels).latent_dist.sample() * scaling + tokens = tokenizer( + caption or "", return_tensors="pt", padding="max_length", + truncation=True, max_length=256, + ).to(device) + embed = text_encoder(**tokens).last_hidden_state + cached.append({"latent": latent.squeeze(0).cpu(), "embed": embed.squeeze(0).cpu()}) + return cached diff --git a/core/src/inline_core/training/models.py b/core/src/inline_core/training/models.py new file mode 100644 index 0000000..853960a --- /dev/null +++ b/core/src/inline_core/training/models.py @@ -0,0 +1,97 @@ +"""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(frozen=True) +class Components: + transformer: Any + 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_components(models_dir: str, base_mode: str, device: str, dtype: Any) -> Components: + """Load the base transformer (grad-enabled, unquantized), VAE, text encoder + tokenizer, and the + flow-match scheduler. In Turbo mode the training adapter is fused into the transformer first.""" + from ..graph.loader_runners import LoraRef + from ..models import loaders + + root = Path(models_dir) + diffusion = _require(root, "diffusion_models", "INLINE_ZIMAGE_MODEL") + vae_file = _require(root, "vae", "INLINE_ZIMAGE_VAE") + encoder_file = _require(root, "text_encoders", "INLINE_ZIMAGE_TEXT_ENCODER") + + adapter = _adapter_path(root, base_mode) + loras: tuple[LoraRef, ...] = (LoraRef(file=adapter, strength=1.0),) if adapter else () + + transformer = loaders.load_diffusion(_ARCH, diffusion, dtype, device=device, loras=loras) + vae = loaders.load_vae(_ARCH, vae_file, dtype, device=device) + text_encoder, tokenizer = loaders.load_text_encoder(_ARCH, encoder_file, dtype, device=device) + scheduler = loaders.load_scheduler(_ARCH) + return Components(transformer, vae, text_encoder, tokenizer, scheduler) 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..c642b82 --- /dev/null +++ b/core/src/inline_core/training/trainer.py @@ -0,0 +1,183 @@ +"""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). + +NOTE (needs a GPU + Z-Image weights to finalize): three Z-Image specifics must be validated against +the runner + the ai-toolkit reference before this trains a good LoRA - + 1. the transformer ``forward`` kwargs / output attribute (``_forward``), + 2. the LoRA ``target_modules`` for ``ZImageTransformer2DModel`` (``_TARGET_MODULES``), + 3. the flow-match timestep scaling + velocity target (``_timesteps`` / the loss). +These are the "study ai-toolkit" items from the plan, not guesses to ship blind. +""" + +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 + +# Attention + FFN projections a Z-Image LoRA typically targets. VALIDATE against the module tree. +_TARGET_MODULES = ["to_q", "to_k", "to_v", "to_out.0", "ff.net.0.proj", "ff.net.2"] + + +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(batch: int, device: Any) -> Any: + import torch + + # Logit-normal sampling of t in (0, 1) - denser near the middle, as flow-match trainers favor. + return torch.sigmoid(torch.randn(batch, device=device)) + + +def _save_checkpoint(accelerator: Any, ckpt_dir: Path, step: int) -> None: + accelerator.save_state(str(ckpt_dir)) + if accelerator.is_main_process: + (ckpt_dir / "step.json").write_text(json.dumps({"step": step}), encoding="utf-8") + + +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, timestep: Any, embeds: Any) -> Any: + """One transformer prediction. VALIDATE the kwargs/return attr against the Z-Image runner.""" + out = transformer( + hidden_states=noisy, timestep=timestep, encoder_hidden_states=embeds, return_dict=True + ) + return out.sample if hasattr(out, "sample") else out[0] + + +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() + + protocol.progress(0, steps, status="loading models") + comps = models.load_components(manifest["modelsDir"], manifest["baseMode"], str(device), dtype) + + protocol.progress(0, steps, status="caching latents") + data = ds.precache(manifest["datasetDir"], comps, str(device), dtype, resolution) + # Precache done: the VAE + text encoder are dead weight for the loop - move them off the GPU so + # the loop only holds the transformer + LoRA + optimizer (the big low-VRAM win). + comps.vae.to("cpu") + comps.text_encoder.to("cpu") + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + transformer = comps.transformer + 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) / "step.json").exists(): + accelerator.load_state(resume_from) + start = _resume_step(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)] + latents = item["latent"].unsqueeze(0).to(device, dtype) + embeds = item["embed"].unsqueeze(0).to(device, dtype) + noise = torch.randn_like(latents) + t = _timesteps(latents.shape[0], device) + t_broadcast = t.view(-1, *([1] * (latents.ndim - 1))) + noisy = (1 - t_broadcast) * latents + t_broadcast * noise + target = noise - latents # rectified-flow velocity target + + with accelerator.accumulate(transformer): + pred = _forward(transformer, noisy, t * 1000.0, embeds) + 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, ckpt_dir, done) + if accelerator.is_main_process: + protocol.checkpoint(str(ckpt_dir)) + + if stop.flagged: + _save_checkpoint(accelerator, 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/src/renderer/components/icons.tsx b/src/renderer/components/icons.tsx index a178b15..8457717 100644 --- a/src/renderer/components/icons.tsx +++ b/src/renderer/components/icons.tsx @@ -189,3 +189,39 @@ 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 ( + + + + ) +} + +/** `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/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..104bf88 --- /dev/null +++ b/src/renderer/store/trainingStore.ts @@ -0,0 +1,217 @@ +/** + * 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, + 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 + 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 + start: (datasetId: string, hyperparams: TrainingHyperparams) => Promise + resume: (runId: string) => Promise + cancel: (runId: string) => Promise + + applyProgress: (e: TrainingProgressEvent) => void + applySample: (e: TrainingSampleEvent) => 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: {}, + 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) return set({ error: res.error }) + set((s) => ({ runs: [res.value, ...s.runs.filter((r) => r.id !== res.value.id)] })) + }, + + 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] }, + })), + + 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)), + })), + + 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.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/Trainer/DatasetItemsGrid.tsx b/src/renderer/views/Trainer/DatasetItemsGrid.tsx new file mode 100644 index 0000000..308efb5 --- /dev/null +++ b/src/renderer/views/Trainer/DatasetItemsGrid.tsx @@ -0,0 +1,113 @@ +/** The dataset editor: a grid of images with per-image captions, plus add + auto-caption controls. */ +import { useMemo, useState } from 'react' +import type { TrainingDatasetItem } from '@shared/types' +import { resolveMedia } from '@/lib/media' +import { studio } from '@/lib/studio' +import { useAssetStore } from '../../store/assetStore' +import { useTrainingStore } from '../../store/trainingStore' +import { ipcErrorMessage } from '../../lib/ipcError' + +function CaptionBox({ + item, + datasetId, +}: { + item: TrainingDatasetItem + datasetId: string +}): React.JSX.Element { + const [text, setText] = useState(item.caption) + const setCaption = useTrainingStore((s) => s.setCaption) + return ( +