diff --git a/README.md b/README.md index af26c96..fbfcf75 100644 --- a/README.md +++ b/README.md @@ -231,9 +231,26 @@ cd core ./webui.sh --install --extra training # Windows: .\webui.bat --install --extra training ``` +**Headless (no UI).** On a rented GPU or in a script, train from a folder of captioned stills or +clips with the same entry point the Trainer tab spawns: + +```bash +cd core +python -m inline_core.training \ + --dataset /path/to/clips \ + --arch minimax-h3 \ + --clip-seconds 1 \ + --models-dir ./models \ + --output ./models/loras/my_h3_clip.safetensors \ + --steps 500 --resolution 512 +``` + +See [Headless training (CLI)](TRAINING.md#headless-training-cli) for the dataset layout, +`metadata.jsonl` import, resume, and the stills-only path. + For a worked example, see [`inlineresearch/skin-lora-krea-2-raw`](https://huggingface.co/inlineresearch/skin-lora-krea-2-raw), a photorealistic skin LoRA trained here on the Krea 2 RAW base from the 26 image and caption pairs published as [`inlineresearch/krea2-skin-lora`](https://huggingface.co/datasets/inlineresearch/krea2-skin-lora). -**[TRAINING.md](TRAINING.md) is the full reference:** [which base to train on](TRAINING.md#architecture-and-base-model-modes) · [measured benchmarks](TRAINING.md#benchmark-results) · [datasets and outputs](TRAINING.md#datasets-and-outputs) · [stop and resume](TRAINING.md#stop-and-resume) · [trigger words](TRAINING.md#trigger-words) · [base precision](TRAINING.md#base-precision) +**[TRAINING.md](TRAINING.md) is the full reference:** [which base to train on](TRAINING.md#architecture-and-base-model-modes) · [training on clips](TRAINING.md#training-on-clips) · [headless CLI](TRAINING.md#headless-training-cli) · [measured benchmarks](TRAINING.md#benchmark-results) · [datasets and outputs](TRAINING.md#datasets-and-outputs) · [stop and resume](TRAINING.md#stop-and-resume) · [trigger words](TRAINING.md#trigger-words) · [base precision](TRAINING.md#base-precision) ## How it works diff --git a/TRAINING.md b/TRAINING.md index b487d98..06fa05c 100644 --- a/TRAINING.md +++ b/TRAINING.md @@ -12,7 +12,8 @@ resolution. **Contents:** [The graph](#the-graph) · [Datasets and outputs](#datasets-and-outputs) · [Stop and resume](#stop-and-resume) · [Trigger words](#trigger-words) · [Architecture and base model modes](#architecture-and-base-model-modes) · [Install](#install) · -[Training on clips](#training-on-clips) · [**Benchmark results**](#benchmark-results) · +[Training on clips](#training-on-clips) · [Headless training (CLI)](#headless-training-cli) · +[**Benchmark results**](#benchmark-results) · [Dataset and adapter options](#dataset-and-adapter-options) · [Base precision](#base-precision) ## The graph @@ -102,6 +103,91 @@ better than the first frame usually does. Write them by hand if you would rather Audio is not trained. H3 generates video and its soundtrack jointly, but the trainer packs zero audio rows, so an adapter changes what a clip looks like and never what it sounds like. +## Headless training (CLI) + +The Trainer tab is optional. The same loop runs from the shell, which is what you want on a rented +GPU, in a Docker box, or when you already have a folder of captioned clips and do not want the SPA. + +**Dataset layout.** A folder of media with captions beside them: + +```text +dataset/ + 0000.mp4 + 0000.txt + 0001.mp4 + 0001.txt + … +``` + +Stills (`.png` / `.jpg` / `.webp`) and clips (`.mp4` / `.mov` / `.webm` / `.mkv`) can mix. If the +folder has a Hugging Face-style `metadata.jsonl` (`file_name` + `text` / `caption`) and no +sidecars, the CLI materialises `NNNN.txt` for you. Put `minimax_h3_fl2va_bf16.safetensors` in +`models/diffusion_models/` first (about 124GB). + +**Train a clip LoRA (motion):** + +```bash +cd core +# after ./webui.sh --install --extra all (or --extra training + runtime) +python -m inline_core.training \ + --dataset /path/to/clips \ + --arch minimax-h3 \ + --clip-seconds 1 \ + --models-dir ./models \ + --output ./models/loras/my_h3_clip.safetensors \ + --steps 500 \ + --resolution 512 \ + --rank 16 +``` + +`--clip-seconds 1` is the practical floor: H3 snaps to a 22-frame grid at 24fps (~0.92s). Longer +clips cost time per step, not VRAM. Omit images-only runs' clip flag if you prefer; any short video +in the set is still snapped to the grid when the arch is H3. + +**Stills only (appearance, no motion):** + +```bash +python -m inline_core.training \ + --dataset /path/to/stills \ + --arch minimax-h3 \ + --models-dir ./models \ + --output ./models/loras/my_h3_still.safetensors \ + --steps 1500 \ + --resolution 512 +``` + +**Useful flags.** `--trigger TOKEN` prepends the token to every caption. `--work-dir DIR` holds the +staged dataset, checkpoints and `manifest.json` (resume with `--resume --work-dir DIR`). +`--dry-run` stages and prints the manifest without loading the GPU. `--gpu-ids 0` pins a card. +Progress is JSON lines on stdout (`type: progress|checkpoint|done|error`), the same protocol the +Trainer node already parses. + +**Weights & Biases (optional).** Set `WANDB_API_KEY` and pass a project (or `WANDB_PROJECT`): + +```bash +export WANDB_API_KEY=… +python -m inline_core.training \ + --dataset /path/to/clips \ + --arch minimax-h3 --clip-seconds 1 \ + --wandb-project scrya-h3-iso-video \ + --wandb-run-name scrya_iso_char_clips \ + --wandb-log-every 5 \ + … +``` + +The run URL is printed on stderr as `wandb: https://wandb.ai/…` and mirrored as a progress status +line. Loss, step, peak VRAM and cache phase strings are logged. Omit `--wandb-project` (and leave +`WANDB_PROJECT` unset) for a fully offline run. + +**Manifest path (what the UI already uses).** You can still hand the subprocess a prebuilt file: + +```bash +python -m inline_core.training /path/to/manifest.json +``` + +That shape is what `studio/training.py` writes; the flag path builds the same object via +`inline_core.training.manifest`. + ## Install If you installed with `--extra all` from [Get Started](README.md#get-started), the trainer is already set up - nothing more to do. To add it to a leaner install, its dependencies (PEFT, 8-bit Adam, the captioner) sit behind the `training` extra: diff --git a/core/src/inline_core/training/__main__.py b/core/src/inline_core/training/__main__.py index 0c10005..0cb974f 100644 --- a/core/src/inline_core/training/__main__.py +++ b/core/src/inline_core/training/__main__.py @@ -1,15 +1,30 @@ -"""Trainer entry point: ``python -m inline_core.training ``. +"""Trainer entry point: headless LoRA training without the Studio UI. -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. +Two invocation styles share the same training loop: + +1. **Manifest** (what the Trainer tab already spawns):: + + python -m inline_core.training /path/to/manifest.json + +2. **Flags** (folder of stills or clips → adapter, no server):: + + python -m inline_core.training \\ + --dataset /data/clips \\ + --arch minimax-h3 \\ + --clip-seconds 1 \\ + --models-dir ./models \\ + --output ./models/loras/my_h3.safetensors \\ + --steps 500 --resolution 512 + +Progress is JSON lines on stdout (``protocol.py``). SIGTERM asks the loop to flush a checkpoint +and stop so the run is resumable with ``--resume`` against the same ``--work-dir``. """ from __future__ import annotations +import argparse import json +import os import sys from pathlib import Path from typing import Any @@ -32,28 +47,233 @@ def _message_for(error: Exception, manifest: dict[str, Any]) -> str: if manifest.get("arch") == "krea2" else "" ) + h3 = ( + " MiniMax H3 peaks around 20.6GB on the caption pass - a 24GB card is the practical floor." + if manifest.get("arch") == "minimax-h3" + else "" + ) return ( - f"Ran out of GPU memory training{at}. Lower the training resolution in the node's Adjust " - f"panel (it drives peak VRAM far more than rank).{lower}{floor}" + f"Ran out of GPU memory training{at}. Lower the training resolution " + f"(it drives peak VRAM far more than rank).{lower}{floor}{h3}" ) -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 +def _parse_gpu_ids(raw: str | None) -> list[int]: + if not raw: + return [] + out: list[int] = [] + for part in raw.split(","): + part = part.strip() + if not part: + continue + out.append(int(part)) + return out + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="python -m inline_core.training", + description=( + "Train a LoRA headlessly. Pass a manifest.json (Trainer-tab shape) or a --dataset " + "folder of captioned stills/clips. MiniMax H3 learns motion when the folder holds " + "clips and --clip-seconds is set (1s is the practical floor)." + ), + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=( + "examples:\n" + " python -m inline_core.training run/manifest.json\n" + " python -m inline_core.training --dataset ./clips --arch minimax-h3 " + "--clip-seconds 1 --models-dir ./models --output ./models/loras/style.safetensors\n" + ), + ) + parser.add_argument( + "manifest", + nargs="?", + help="Existing run manifest.json (same shape the Trainer tab writes).", + ) + parser.add_argument( + "--dataset", + type=Path, + help="Folder of NNNN.png/mp4 + NNNN.txt pairs (or metadata.jsonl captions).", + ) + parser.add_argument( + "--models-dir", + type=Path, + default=None, + help="Models root (diffusion_models/, vae/, …). Default: INLINE_MODELS_DIR or ./models.", + ) + parser.add_argument( + "--output", + type=Path, + default=None, + help="Where to write the finished .safetensors LoRA.", + ) + parser.add_argument( + "--work-dir", + type=Path, + default=None, + help="Working directory for the staged dataset, checkpoints, and manifest.", + ) + parser.add_argument( + "--arch", + choices=("z-image", "krea2", "flux2", "minimax-h3"), + default="minimax-h3", + help="Training architecture. Default: minimax-h3 (the video model).", + ) + parser.add_argument( + "--base-mode", + default=None, + help="Base checkpoint mode (H3: raw / FL2VA). Defaults per architecture.", + ) + parser.add_argument( + "--clip-seconds", + type=float, + default=None, + help=( + "H3 only: train clips as motion. 1s is the practical floor (22 frames at 24fps). " + "Omit for stills-only intent; short clips are still snapped to the grid." + ), + ) + parser.add_argument("--rank", type=int, default=16) + parser.add_argument("--alpha", type=int, default=None, help="Defaults to --rank.") + parser.add_argument("--learning-rate", type=float, default=1e-4) + parser.add_argument("--steps", type=int, default=500) + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--resolution", type=int, default=512) + parser.add_argument("--save-every", type=int, default=250) + parser.add_argument( + "--base-quant", + choices=("auto", "none", "nf4"), + default="auto", + help="H3 is 4-bit only; this is accepted and forced to nf4 internally.", + ) + parser.add_argument("--offload", choices=("auto", "on", "off"), default="auto") + parser.add_argument("--lora-scope", choices=("full", "attention"), default="full") + parser.add_argument("--caption-dropout", type=float, default=0.05) + parser.add_argument( + "--flip", + action="store_true", + help="Mirror every still (exact pixel flip, not latent flip).", + ) + parser.add_argument("--trigger", default="", help="Prepended to every caption.") + parser.add_argument( + "--output-name", + default="", + help="Stem used when --output is omitted (writes models/loras/.safetensors).", + ) + parser.add_argument( + "--gpu-ids", + default="", + help="Comma-separated GPU indices (sets CUDA_VISIBLE_DEVICES). Multi-GPU uses accelerate.", + ) + parser.add_argument( + "--resume", + action="store_true", + help="Continue from checkpoints in --work-dir (or the manifest's checkpointDir).", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Stage the dataset and write the manifest, then print its path and exit.", + ) + parser.add_argument( + "--wandb-project", + default="", + help="Weights & Biases project (or set WANDB_PROJECT). Empty = no wandb.", + ) + parser.add_argument( + "--wandb-run-name", + default="", + help="Weights & Biases run name (or set WANDB_RUN_NAME).", + ) + parser.add_argument( + "--wandb-log-every", + type=int, + default=1, + help="Log train/loss to wandb every N steps (default 1).", + ) + return parser + + +def _default_models_dir() -> Path: + env = os.environ.get("INLINE_MODELS_DIR") + return Path(env).expanduser() if env else Path("models") + + +def _resolve_manifest(args: argparse.Namespace) -> dict[str, Any]: + """Load a manifest from the positional path or build one from flags.""" + if args.manifest: + path = Path(args.manifest) + try: + return json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise SystemExit(f"Could not read manifest: {exc}") from exc + if not args.dataset: + raise SystemExit( + "Pass a manifest.json path or --dataset pointing at a folder of stills/clips." + ) + + from .manifest import prepare_run + + models = args.models_dir or _default_models_dir() + # Default H3 clip length when the arch is H3 and the user did not pass the flag: 1s is the + # practical floor and matches the Trainer Adjust panel default. + clip_seconds = args.clip_seconds + if args.arch == "minimax-h3" and clip_seconds is None: + clip_seconds = 1.0 + + path, manifest = prepare_run( + dataset=args.dataset, + models_dir=models, + output=args.output, + work_dir=args.work_dir, + arch=args.arch, + base_mode=args.base_mode, + trigger=args.trigger, + resume=args.resume, + output_name=args.output_name, + rank=args.rank, + alpha=args.alpha, + learning_rate=args.learning_rate, + steps=args.steps, + batch_size=args.batch_size, + resolution=args.resolution, + save_every=args.save_every, + base_quant=args.base_quant, + offload=args.offload, + lora_scope=args.lora_scope, + caption_dropout=args.caption_dropout, + flip_augment=args.flip, + clip_seconds=clip_seconds, + gpu_ids=_parse_gpu_ids(args.gpu_ids), + ) + # WandB knobs live on the hyperparams block (and/or env) so a prebuilt manifest can carry them. + if args.wandb_project: + manifest["hyperparams"]["wandbProject"] = args.wandb_project + os.environ.setdefault("WANDB_PROJECT", args.wandb_project) + if args.wandb_run_name: + manifest["hyperparams"]["wandbRunName"] = args.wandb_run_name + os.environ.setdefault("WANDB_RUN_NAME", args.wandb_run_name) + if args.wandb_log_every: + manifest["hyperparams"]["wandbLogEvery"] = int(args.wandb_log_every) + # Surface the path on stderr so stdout stays the pure progress protocol. + print(f"manifest: {path}", file=sys.stderr) + return manifest + + +def run_training(manifest: dict[str, Any]) -> int: 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 + gpu_ids = [int(g) for g in (manifest.get("gpuIds") or [])] + if gpu_ids: + os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(str(g) for g in gpu_ids) + os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + try: output = train(manifest) except KeyboardInterrupt: @@ -70,5 +290,47 @@ def main(argv: list[str]) -> int: return 0 +def main(argv: list[str] | None = None) -> int: + argv = list(sys.argv[1:] if argv is None else argv) + + # Back-compat: a single non-flag argument is the manifest path the Studio orchestrator passes. + if len(argv) == 1 and not argv[0].startswith("-"): + 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 + return run_training(manifest) + + if not argv: + protocol.error( + "No manifest path given. Pass manifest.json or use --dataset /path (see --help)." + ) + return 2 + + parser = build_parser() + try: + args = parser.parse_args(argv) + except SystemExit as exc: + code = exc.code + return 0 if code in (0, None) else (code if isinstance(code, int) else 2) + + try: + manifest = _resolve_manifest(args) + except SystemExit as exc: + message = exc.code if isinstance(exc.code, str) else str(exc) + if message and message not in ("0", "1", "2"): + protocol.error(message) + return 2 + except Exception as exc: # noqa: BLE001 - staging failures before torch loads + protocol.error(str(exc)) + return 2 + + if args.dry_run: + print(json.dumps(manifest, indent=2)) + return 0 + return run_training(manifest) + + if __name__ == "__main__": - raise SystemExit(main(sys.argv[1:])) + raise SystemExit(main()) diff --git a/core/src/inline_core/training/manifest.py b/core/src/inline_core/training/manifest.py new file mode 100644 index 0000000..9f93fc6 --- /dev/null +++ b/core/src/inline_core/training/manifest.py @@ -0,0 +1,266 @@ +"""Build a trainer run manifest without the Studio UI. + +The Trainer tab writes the same shape from ``studio/training.py::_prepare``. Headless runs +(``python -m inline_core.training --dataset …``) use this module so a folder of stills or clips +trains through the same entry point, without starting the server or the SPA. +""" + +from __future__ import annotations + +import json +import re +import shutil +import uuid +from pathlib import Path +from typing import Any + +from . import dataset as ds + +_ARCHS = ("z-image", "krea2", "flux2", "minimax-h3") +_DEFAULT_BASE = { + "z-image": "deturbo", + "krea2": "raw", + "flux2": "raw", + "minimax-h3": "raw", +} + + +def _safe(name: str) -> str: + return re.sub(r"[^\w.-]+", "_", name).strip("_") or "lora" + + +def default_base_mode(arch: str) -> str: + if arch not in _DEFAULT_BASE: + raise ValueError(f"Unknown training architecture {arch!r}. Choose one of {_ARCHS}.") + return _DEFAULT_BASE[arch] + + +def load_metadata_jsonl(folder: Path) -> dict[str, str]: + """Optional Hugging Face-style captions: one JSON object per line with a file name and text. + + Looks for ``file_name`` / ``file`` / ``path`` and ``text`` / ``caption`` / ``prompt``. Missing + keys are skipped; sidecars always win when both exist. + """ + path = folder / "metadata.jsonl" + if not path.is_file(): + return {} + out: dict[str, str] = {} + for line in path.read_text(encoding="utf-8").splitlines(): + line = line.strip() + if not line: + continue + try: + row = json.loads(line) + except json.JSONDecodeError: + continue + if not isinstance(row, dict): + continue + name = row.get("file_name") or row.get("file") or row.get("path") + caption = row.get("text") or row.get("caption") or row.get("prompt") + if not isinstance(name, str) or not isinstance(caption, str): + continue + out[Path(name).name] = caption.strip() + return out + + +def list_media(folder: Path, *, allow_video: bool) -> list[Path]: + """Media files in a dataset folder, images always, clips only when the arch can train them.""" + suffixes = ds._IMAGE_SUFFIXES + (ds._VIDEO_SUFFIXES if allow_video else ()) + return sorted( + p for p in folder.iterdir() if p.is_file() and p.suffix.lower() in suffixes + ) + + +def stage_dataset( + source: Path, + dest: Path, + *, + trigger: str = "", + allow_video: bool = True, +) -> int: + """Copy media + captions into ``dest`` as ``NNNN.ext`` / ``NNNN.txt`` pairs. + + Sidecar ``.txt`` captions win. When a file has none, ``metadata.jsonl`` in the source folder is + used if present. ``trigger`` is prepended the same way the Studio export does. + """ + source = source.expanduser().resolve() + if not source.is_dir(): + raise ValueError(f"Dataset folder not found: {source}") + media = list_media(source, allow_video=allow_video) + if not media: + kinds = "images or clips" if allow_video else "images" + raise ValueError(f"No {kinds} in {source}") + + dest.mkdir(parents=True, exist_ok=True) + captions = load_metadata_jsonl(source) + trigger = trigger.strip() + for index, path in enumerate(media): + out = dest / f"{index:04d}{path.suffix.lower()}" + shutil.copyfile(path, out) + sidecar = path.with_suffix(".txt") + if sidecar.is_file(): + text = sidecar.read_text(encoding="utf-8").strip() + else: + text = captions.get(path.name, "").strip() + caption = ", ".join(part for part in (trigger, text) if part) + out.with_suffix(".txt").write_text(caption, encoding="utf-8") + return len(media) + + +def build_hyperparams( + *, + arch: str, + base_mode: str | None = None, + rank: int = 16, + alpha: int | None = None, + learning_rate: float = 1e-4, + steps: int = 500, + batch_size: int = 1, + resolution: int = 512, + save_every: int = 250, + base_quant: str = "auto", + offload: str = "auto", + lora_scope: str = "full", + caption_dropout: float = 0.05, + flip_augment: bool = False, + clip_seconds: float | None = None, + output_name: str = "", + gpu_ids: list[int] | None = None, + wandb_project: str = "", + wandb_run_name: str = "", + wandb_log_every: int = 1, +) -> dict[str, Any]: + """Hyperparams block matching what the Trainer Adjust panel posts.""" + if arch not in _ARCHS: + raise ValueError(f"Unknown training architecture {arch!r}. Choose one of {_ARCHS}.") + hp: dict[str, Any] = { + "arch": arch, + "baseMode": base_mode or default_base_mode(arch), + "baseQuant": base_quant, + "offload": offload, + "loraScope": lora_scope, + "captionDropout": float(caption_dropout), + "flipAugment": bool(flip_augment), + "rank": int(rank), + "alpha": int(alpha if alpha is not None else rank), + "learningRate": float(learning_rate), + "batchSize": int(batch_size), + "steps": int(steps), + "saveEvery": int(save_every), + "resolution": int(resolution), + "outputName": output_name, + "gpuIds": list(gpu_ids or []), + } + # H3 only: omit for stills-only intent when the caller leaves it unset. The trainer still + # snaps an unset clip length to the 22-frame floor for any video files present. + if arch == "minimax-h3" and clip_seconds is not None: + hp["clipSeconds"] = float(clip_seconds) + if wandb_project: + hp["wandbProject"] = wandb_project + if wandb_run_name: + hp["wandbRunName"] = wandb_run_name + if wandb_log_every and wandb_log_every != 1: + hp["wandbLogEvery"] = int(wandb_log_every) + return hp + + +def build_manifest( + *, + dataset_dir: Path | str, + models_dir: Path | str, + output_path: Path | str, + work_dir: Path | str, + hyperparams: dict[str, Any], + run_id: str | None = None, + resume: bool = False, + trigger_word: str = "", +) -> dict[str, Any]: + """The JSON object ``python -m inline_core.training `` consumes.""" + work = Path(work_dir) + checkpoints = work / "checkpoints" + checkpoints.mkdir(parents=True, exist_ok=True) + rid = run_id or uuid.uuid4().hex[:12] + return { + "runId": rid, + "workingDir": str(work), + "datasetDir": str(Path(dataset_dir)), + "checkpointDir": str(checkpoints), + "outputPath": str(Path(output_path)), + "resumeFrom": str(checkpoints) if resume else None, + "modelsDir": str(Path(models_dir)), + "arch": hyperparams.get("arch") or "z-image", + "baseMode": hyperparams["baseMode"], + "triggerWord": trigger_word, + "hyperparams": hyperparams, + "gpuIds": hyperparams.get("gpuIds") or [], + } + + +def write_manifest(manifest: dict[str, Any], path: Path | str) -> Path: + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(manifest, indent=2), encoding="utf-8") + return path + + +def prepare_run( + *, + dataset: Path | str, + models_dir: Path | str, + output: Path | str | None = None, + work_dir: Path | str | None = None, + arch: str = "minimax-h3", + base_mode: str | None = None, + trigger: str = "", + resume: bool = False, + output_name: str = "", + **hp_kwargs: Any, +) -> tuple[Path, dict[str, Any]]: + """Stage the dataset if needed, write the manifest, return ``(manifest_path, manifest)``. + + Always stages into ``work_dir/dataset`` so captions can pick up a trigger word and + ``metadata.jsonl`` without mutating the user's folder. + """ + source = Path(dataset).expanduser().resolve() + rid = uuid.uuid4().hex[:12] + work = Path(work_dir).expanduser().resolve() if work_dir else Path("training_runs") / rid + work.mkdir(parents=True, exist_ok=True) + + hp = build_hyperparams( + arch=arch, + base_mode=base_mode, + output_name=output_name, + **hp_kwargs, + ) + staged = work / "dataset" + if staged.exists(): + shutil.rmtree(staged) + count = stage_dataset( + source, + staged, + trigger=trigger, + allow_video=arch == "minimax-h3", + ) + if count == 0: + raise ValueError(f"No training items staged from {source}") + + models = Path(models_dir).expanduser().resolve() + if output is not None: + out = Path(output).expanduser().resolve() + else: + stem = _safe(output_name or f"{arch}-{rid}") + out = models / "loras" / f"{stem}.safetensors" + out.parent.mkdir(parents=True, exist_ok=True) + + manifest = build_manifest( + dataset_dir=staged, + models_dir=models, + output_path=out, + work_dir=work, + hyperparams=hp, + run_id=rid, + resume=resume, + trigger_word=trigger, + ) + path = write_manifest(manifest, work / "manifest.json") + return path, manifest diff --git a/core/src/inline_core/training/trainer.py b/core/src/inline_core/training/trainer.py index a0352f6..0b5a696 100644 --- a/core/src/inline_core/training/trainer.py +++ b/core/src/inline_core/training/trainer.py @@ -150,6 +150,59 @@ def _to_device(item: dict[str, Any], device: Any, dtype: Any) -> dict[str, Any]: } +def _wandb_run(manifest: dict[str, Any], steps: int) -> Any: + """Optional Weights & Biases run for headless / Vast. Off unless WANDB_PROJECT is set + (or manifest hyperparams carry wandbProject) and the wandb package is installed.""" + import os + + hp = manifest.get("hyperparams") or {} + project = ( + os.environ.get("WANDB_PROJECT") + or hp.get("wandbProject") + or "" + ).strip() + if not project: + return None + try: + import wandb + except ImportError: + protocol.progress(0, steps, status="wandb requested but package missing; continuing offline") + return None + name = ( + os.environ.get("WANDB_RUN_NAME") + or hp.get("wandbRunName") + or hp.get("outputName") + or manifest.get("runId") + or "inline-lora" + ) + config = { + "arch": manifest.get("arch"), + "baseMode": manifest.get("baseMode"), + "steps": steps, + "rank": hp.get("rank"), + "alpha": hp.get("alpha"), + "learningRate": hp.get("learningRate"), + "resolution": hp.get("resolution"), + "clipSeconds": hp.get("clipSeconds"), + "batchSize": hp.get("batchSize"), + "loraScope": hp.get("loraScope"), + "datasetDir": manifest.get("datasetDir"), + } + run = wandb.init( + project=project, + name=str(name), + config=config, + dir=str(Path(manifest.get("workingDir") or ".")), + resume="allow", + ) + # JSON protocol stays machine-readable; surface the human URL on stderr + a progress status. + url = getattr(run, "url", None) or "" + if url: + print(f"wandb: {url}", flush=True) + protocol.progress(0, steps, status=f"wandb {url}") + return run + + def train(manifest: dict[str, Any]) -> str | None: import torch from accelerate import Accelerator @@ -161,6 +214,7 @@ def train(manifest: dict[str, Any]) -> str | None: hp = manifest["hyperparams"] steps = int(hp["steps"]) save_every = max(1, int(hp.get("saveEvery", 250))) + log_every = max(1, int(hp.get("wandbLogEvery") or hp.get("logEvery") or 1)) resolution = int(hp.get("resolution", 1024)) ckpt_dir = Path(manifest["checkpointDir"]) ckpt_dir.mkdir(parents=True, exist_ok=True) @@ -169,6 +223,8 @@ def train(manifest: dict[str, Any]) -> str | None: device = accelerator.device dtype = models.compute_dtype() + wb = _wandb_run(manifest, steps) if accelerator.is_main_process else None + # Two phases, never overlapping: encoders -> precache -> free, THEN the transformer. Held # together they add the text encoder's several GB to the base; apart, peak is just the base. # Before the precache, never after: this costs milliseconds and precaching costs twenty @@ -180,12 +236,21 @@ def train(manifest: dict[str, Any]) -> str | None: # Precache is minutes of silence on a large dataset, so its phases are reported as progress # statuses. The orchestrator turns each new status into a log line, which is the only channel # that reaches the UI: this subprocess installs no logging handler. + def _cache_status(text: str) -> None: + protocol.progress(0, steps, status=text) + if wb is not None: + wb.log({"status": text}) + data, unconditional, shift = cache.build( manifest["datasetDir"], manifest["modelsDir"], arch.key, str(device), dtype, resolution, flip=bool(hp.get("flipAugment")), dropout=dropout, clip_frames=archs.clip_frames(arch, hp.get("clipSeconds")), - on_status=lambda text: protocol.progress(0, steps, status=text), + on_status=_cache_status, ) + if wb is not None: + wb.log({"dataset/items": len(data), "train/clip_frames": archs.clip_frames( + arch, hp.get("clipSeconds") + )}) quant = models.resolve_quant( str(hp.get("baseQuant") or "auto"), @@ -264,19 +329,37 @@ def train(manifest: dict[str, Any]) -> str | None: optimizer.zero_grad() done = step + 1 - protocol.progress( - done, steps, loss=float(loss.detach().item()), status="training", vram=_peak_vram_gb() - ) + loss_val = float(loss.detach().item()) + vram = _peak_vram_gb() + protocol.progress(done, steps, loss=loss_val, status="training", vram=vram) + if wb is not None and (done % log_every == 0 or done == steps or done == 1): + payload: dict[str, Any] = { + "train/loss": loss_val, + "train/step": done, + "train/epoch_frac": done / max(len(data), 1), + } + if vram is not None: + payload["system/peak_vram_gb"] = vram + wb.log(payload, step=done) if done % save_every == 0 or done == steps: _save_checkpoint(accelerator, transformer, optimizer, ckpt_dir, done) if accelerator.is_main_process: protocol.checkpoint(str(ckpt_dir)) + if wb is not None: + wb.log({"train/checkpoint_step": done}, step=done) if stop.flagged: _save_checkpoint(accelerator, transformer, optimizer, ckpt_dir, step) + if wb is not None: + wb.log({"train/status": "interrupted", "train/step": step}) + wb.finish() 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"]) + if wb is not None: + wb.log({"train/status": "done", "train/step": steps}) + # Finish so the run URL is complete even if the process is killed on Vast teardown. + wb.finish() return manifest["outputPath"] diff --git a/core/tests/test_training_cli.py b/core/tests/test_training_cli.py new file mode 100644 index 0000000..7690896 --- /dev/null +++ b/core/tests/test_training_cli.py @@ -0,0 +1,203 @@ +"""Headless trainer CLI: manifest building and flag parsing, without a GPU.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from inline_core.training import __main__ as entry +from inline_core.training import manifest as mf + + +def _touch_media(folder: Path, name: str, caption: str | None = None) -> Path: + path = folder / name + path.write_bytes(b"not-a-real-image") + if caption is not None: + path.with_suffix(".txt").write_text(caption, encoding="utf-8") + return path + + +def test_stage_dataset_copies_sidecars_and_trigger(tmp_path: Path) -> None: + src = tmp_path / "src" + src.mkdir() + _touch_media(src, "hero.png", "a knight") + _touch_media(src, "walk.mp4", "walking left") + + dest = tmp_path / "staged" + count = mf.stage_dataset(src, dest, trigger="sks", allow_video=True) + + assert count == 2 + assert (dest / "0000.png").is_file() + assert (dest / "0001.mp4").is_file() + assert (dest / "0000.txt").read_text(encoding="utf-8") == "sks, a knight" + assert (dest / "0001.txt").read_text(encoding="utf-8") == "sks, walking left" + + +def test_stage_dataset_reads_metadata_jsonl_when_sidecars_missing(tmp_path: Path) -> None: + src = tmp_path / "src" + src.mkdir() + _touch_media(src, "clip_a.mp4") + _touch_media(src, "clip_b.mp4") + (src / "metadata.jsonl").write_text( + "\n".join( + [ + json.dumps({"file_name": "clip_a.mp4", "text": "pixel walk cycle"}), + json.dumps({"file": "clip_b.mp4", "caption": "pixel idle"}), + ] + ), + encoding="utf-8", + ) + + dest = tmp_path / "staged" + mf.stage_dataset(src, dest, allow_video=True) + + assert (dest / "0000.txt").read_text(encoding="utf-8") == "pixel walk cycle" + assert (dest / "0001.txt").read_text(encoding="utf-8") == "pixel idle" + + +def test_stage_dataset_sidecar_wins_over_jsonl(tmp_path: Path) -> None: + src = tmp_path / "src" + src.mkdir() + _touch_media(src, "a.mp4", "from sidecar") + (src / "metadata.jsonl").write_text( + json.dumps({"file_name": "a.mp4", "text": "from jsonl"}), encoding="utf-8" + ) + + dest = tmp_path / "staged" + mf.stage_dataset(src, dest, allow_video=True) + assert (dest / "0000.txt").read_text(encoding="utf-8") == "from sidecar" + + +def test_image_arch_skips_clips(tmp_path: Path) -> None: + src = tmp_path / "src" + src.mkdir() + _touch_media(src, "still.jpg", "face") + _touch_media(src, "motion.mp4", "walk") + + dest = tmp_path / "staged" + count = mf.stage_dataset(src, dest, allow_video=False) + assert count == 1 + assert (dest / "0000.jpg").is_file() + assert not (dest / "0001.mp4").exists() + + +def test_build_hyperparams_includes_clip_seconds_for_h3() -> None: + hp = mf.build_hyperparams(arch="minimax-h3", clip_seconds=1.0, steps=100) + assert hp["arch"] == "minimax-h3" + assert hp["baseMode"] == "raw" + assert hp["clipSeconds"] == 1.0 + assert hp["steps"] == 100 + + +def test_build_hyperparams_omits_clip_seconds_for_image_archs() -> None: + hp = mf.build_hyperparams(arch="z-image", clip_seconds=2.0) + assert "clipSeconds" not in hp + + +def test_prepare_run_writes_manifest(tmp_path: Path) -> None: + src = tmp_path / "src" + src.mkdir() + _touch_media(src, "0000.mp4", "pixel art walk") + models = tmp_path / "models" + models.mkdir() + work = tmp_path / "work" + out = tmp_path / "out" / "lora.safetensors" + + path, manifest = mf.prepare_run( + dataset=src, + models_dir=models, + output=out, + work_dir=work, + arch="minimax-h3", + clip_seconds=1.0, + steps=12, + resolution=512, + trigger="pxart", + ) + + assert path == work / "manifest.json" + assert path.is_file() + assert manifest["arch"] == "minimax-h3" + assert manifest["hyperparams"]["clipSeconds"] == 1.0 + assert manifest["outputPath"] == str(out.resolve()) + assert Path(manifest["datasetDir"]).is_dir() + caption = next(Path(manifest["datasetDir"]).glob("*.txt")).read_text(encoding="utf-8") + assert caption == "pxart, pixel art walk" + + +def test_cli_dry_run_prints_manifest(tmp_path: Path, capsys: pytest.CaptureFixture[str]) -> None: + src = tmp_path / "src" + src.mkdir() + _touch_media(src, "a.png", "still") + models = tmp_path / "models" + models.mkdir() + work = tmp_path / "work" + + code = entry.main( + [ + "--dataset", + str(src), + "--models-dir", + str(models), + "--work-dir", + str(work), + "--arch", + "minimax-h3", + "--clip-seconds", + "1", + "--steps", + "8", + "--dry-run", + ] + ) + + assert code == 0 + printed = json.loads(capsys.readouterr().out) + assert printed["arch"] == "minimax-h3" + assert printed["hyperparams"]["clipSeconds"] == 1.0 + assert printed["hyperparams"]["steps"] == 8 + assert (work / "manifest.json").is_file() + + +def test_cli_manifest_path_still_works(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """Studio orchestrator path: a single manifest.json argument must not require flags.""" + manifest = { + "arch": "minimax-h3", + "baseMode": "raw", + "hyperparams": {"resolution": 512, "steps": 1}, + "outputPath": str(tmp_path / "out.safetensors"), + } + path = tmp_path / "manifest.json" + path.write_text(json.dumps(manifest), encoding="utf-8") + + called: dict[str, object] = {} + + def fake_run(m: dict) -> int: + called["m"] = m + return 0 + + monkeypatch.setattr(entry, "run_training", fake_run) + code = entry.main([str(path)]) + assert code == 0 + assert called["m"]["arch"] == "minimax-h3" + + +def test_cli_missing_dataset_errors(capsys: pytest.CaptureFixture[str]) -> None: + code = entry.main(["--arch", "minimax-h3", "--dry-run"]) + assert code == 2 + err = capsys.readouterr() + # protocol.error goes to stdout as JSON + lines = [line for line in err.out.splitlines() if line.startswith("{")] + assert lines + payload = json.loads(lines[-1]) + assert payload["type"] == "error" + assert "dataset" in payload["message"].lower() + + +def test_cli_empty_args_errors(capsys: pytest.CaptureFixture[str]) -> None: + code = entry.main([]) + assert code == 2 + payload = json.loads(capsys.readouterr().out.strip().splitlines()[-1]) + assert payload["type"] == "error"