Merge remote-tracking branch 'origin/diffusion-train-tab-2' into fold-integration
This commit is contained in:
Daniel Han
commit
b52e7a5cc2
12 changed files with 1433 additions and 206 deletions
|
|
@ -31,7 +31,7 @@ import os
|
|||
import random
|
||||
import time
|
||||
from contextlib import nullcontext
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, replace
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Optional
|
||||
|
||||
|
|
@ -52,7 +52,9 @@ from core.training.diffusion_train_common import (
|
|||
_restore_perf_flags,
|
||||
discover_image_caption_pairs,
|
||||
has_functional_torchao,
|
||||
PermutationBatchSampler,
|
||||
repo_is_prequantized,
|
||||
resolve_train_steps,
|
||||
)
|
||||
|
||||
# Per-family LoRA target modules (attention projections). FLUX / Qwen double-stream blocks
|
||||
|
|
@ -1038,6 +1040,10 @@ def run_dit_lora_training(
|
|||
pairs = discover_image_caption_pairs(
|
||||
cfg.data_dir, instance_prompt = cfg.instance_prompt, caption_column = cfg.caption_column
|
||||
)
|
||||
# Resolve num_epochs -> a concrete train_steps now that the dataset size is known, and
|
||||
# rebind cfg so every downstream read (scheduler length, the loop range, progress
|
||||
# total_steps, steps_run) sees the same resolved value.
|
||||
cfg = replace(cfg, train_steps = resolve_train_steps(cfg, len(pairs)), num_epochs = 0)
|
||||
_emit(on_event, "model_load_started", num_images = len(pairs))
|
||||
if _check_stop():
|
||||
out_dir = Path(cfg.output_dir).expanduser()
|
||||
|
|
@ -1199,6 +1205,10 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
|
|||
transformer.train()
|
||||
n_images = len(image_paths)
|
||||
batch_size = cfg.train_batch_size
|
||||
# Permutation-cycle index sampler (shared with the SDXL trainer): visits every image once
|
||||
# per cycle before repeating, so a short run covers the whole dataset instead of the old
|
||||
# with-replacement draw. Uses the loop's own rng to stay seed-deterministic.
|
||||
index_sampler = PermutationBatchSampler(n_images, rng)
|
||||
stopped = False
|
||||
running_loss = 0.0
|
||||
peak_gb = 0.0
|
||||
|
|
@ -1218,7 +1228,7 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
|
|||
optimizer.zero_grad(set_to_none = True)
|
||||
step_loss = 0.0
|
||||
for _ in range(cfg.gradient_accumulation_steps):
|
||||
idxs = [rng.randrange(n_images) for _ in range(batch_size)]
|
||||
idxs = index_sampler.next_batch(batch_size)
|
||||
if latent_cache is not None:
|
||||
latents = _sample_cached_latents(
|
||||
latent_cache, idxs, variant_rng, device, weight_dtype
|
||||
|
|
@ -1254,10 +1264,10 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
|
|||
(loss / cfg.gradient_accumulation_steps).backward()
|
||||
step_loss += float(loss.detach()) / cfg.gradient_accumulation_steps
|
||||
|
||||
grad_norm: Optional[float] = None
|
||||
grad_norm = None
|
||||
if cfg.max_grad_norm and cfg.max_grad_norm > 0:
|
||||
# clip_grad_norm_ returns the PRE-clip total norm: the signal the grad-norm
|
||||
# chart wants (spikes stay visible even when clipping flattens the update).
|
||||
# clip_grad_norm_ returns the total PRE-clip norm: the health signal the UI
|
||||
# charts (an exploding norm shows up here even while the clip caps the update).
|
||||
grad_norm = float(torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm))
|
||||
optimizer.step()
|
||||
lr_sched.step()
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ import gc
|
|||
import os
|
||||
import random
|
||||
import time
|
||||
from dataclasses import replace
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
|
|
@ -61,6 +62,8 @@ from core.training.diffusion_train_common import ( # noqa: F401
|
|||
_restore_perf_flags,
|
||||
discover_image_caption_pairs,
|
||||
get_trainer,
|
||||
PermutationBatchSampler,
|
||||
resolve_train_steps,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -331,6 +334,10 @@ def run_diffusion_lora_training(
|
|||
pairs = discover_image_caption_pairs(
|
||||
cfg.data_dir, instance_prompt = cfg.instance_prompt, caption_column = cfg.caption_column
|
||||
)
|
||||
# Resolve num_epochs -> a concrete train_steps now that the dataset size is known, and
|
||||
# rebind cfg so every downstream read (scheduler length, the loop range, progress
|
||||
# total_steps, steps_run) sees the same resolved value.
|
||||
cfg = replace(cfg, train_steps = resolve_train_steps(cfg, len(pairs)), num_epochs = 0)
|
||||
_emit(on_event, "model_load_started", num_images = len(pairs))
|
||||
|
||||
# Honour a stop requested before the (potentially large / slow) base model loads, the
|
||||
|
|
@ -460,8 +467,14 @@ def run_diffusion_lora_training(
|
|||
|
||||
_emit(on_event, "model_load_completed")
|
||||
|
||||
# Permutation-cycle index sampler (shared with the DiT trainer): each dataset image is
|
||||
# visited once per cycle before any repeat, so a short run does not leave part of a
|
||||
# small dataset unseen. Draws from the loop's own rng so the sequence stays
|
||||
# seed-deterministic.
|
||||
index_sampler = PermutationBatchSampler(len(pairs), rng)
|
||||
|
||||
def _next_batch() -> tuple[list[int], list[str], list[str]]:
|
||||
idx = rng.sample(range(len(pairs)), k = min(cfg.train_batch_size, len(pairs)))
|
||||
idx = index_sampler.next_batch(min(cfg.train_batch_size, len(pairs)))
|
||||
chosen = [pairs[i] for i in idx]
|
||||
return idx, [c[0] for c in chosen], [c[1] for c in chosen]
|
||||
|
||||
|
|
@ -545,9 +558,9 @@ def run_diffusion_lora_training(
|
|||
|
||||
# max_grad_norm <= 0 means "disable clipping" (the Studio payload sends 0.0 for that);
|
||||
# passing 0.0 to clip_grad_norm_ would scale every gradient to zero (no learning).
|
||||
grad_norm: Optional[float] = None
|
||||
grad_norm = None
|
||||
if cfg.max_grad_norm and cfg.max_grad_norm > 0:
|
||||
# clip_grad_norm_ returns the PRE-clip total norm (the grad-norm chart signal).
|
||||
# The returned value is the total PRE-clip norm, reported to the UI chart.
|
||||
grad_norm = float(torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm))
|
||||
optimizer.step()
|
||||
lr_sched.step()
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ actual training loop; this module only routes a request to the right one.
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
|
|
@ -366,6 +367,9 @@ class DiffusionLoraConfig:
|
|||
instance_prompt: Optional[str] = None
|
||||
resolution: int = 1024
|
||||
train_steps: int = 500
|
||||
# 0 = disabled (train for train_steps). > 0 overrides train_steps with a run length of
|
||||
# num_epochs full passes over the dataset, in optimizer steps (see resolve_train_steps).
|
||||
num_epochs: int = 0
|
||||
learning_rate: float = 1e-4
|
||||
train_batch_size: int = 1
|
||||
gradient_accumulation_steps: int = 1
|
||||
|
|
@ -419,6 +423,8 @@ class DiffusionLoraConfig:
|
|||
resolved_family = resolve_trainable_family(self.base_model, self.model_family)
|
||||
if self.train_steps < 1:
|
||||
raise ValueError("train_steps must be >= 1")
|
||||
if not 0 <= int(self.num_epochs) <= 1000:
|
||||
raise ValueError("num_epochs must be between 0 and 1000 (0 uses train_steps)")
|
||||
if self.train_batch_size < 1:
|
||||
raise ValueError("train_batch_size must be >= 1")
|
||||
if self.gradient_accumulation_steps < 1:
|
||||
|
|
@ -494,6 +500,7 @@ class DiffusionLoraConfig:
|
|||
lora_target_modules = targets,
|
||||
max_grad_norm = float(self.max_grad_norm),
|
||||
hf_token = token or None,
|
||||
num_epochs = int(self.num_epochs),
|
||||
cache_variants = int(self.cache_variants),
|
||||
compile_transformer = compile_transformer,
|
||||
base_precision = base_precision,
|
||||
|
|
@ -501,6 +508,59 @@ class DiffusionLoraConfig:
|
|||
)
|
||||
|
||||
|
||||
def resolve_train_steps(cfg: "DiffusionLoraConfig", n_images: int) -> int:
|
||||
"""The effective optimizer-step count for a run. When ``cfg.num_epochs`` is set (> 0),
|
||||
one epoch is one full pass over the dataset in optimizer steps -- ceil(N / (batch x
|
||||
grad_accum)) steps -- so the run is ``num_epochs`` such passes, capped at 100000. With
|
||||
``num_epochs == 0`` the explicit ``cfg.train_steps`` is used unchanged."""
|
||||
if cfg.num_epochs > 0:
|
||||
per_step = max(1, cfg.train_batch_size * cfg.gradient_accumulation_steps)
|
||||
steps_per_epoch = max(1, math.ceil(n_images / per_step))
|
||||
return min(100000, cfg.num_epochs * steps_per_epoch)
|
||||
return cfg.train_steps
|
||||
|
||||
|
||||
class PermutationBatchSampler:
|
||||
"""Yields batch indices as consecutive slices of a reshuffled permutation of
|
||||
``range(n)``, so every index is visited exactly once per cycle before any repeats --
|
||||
an epoch-style full pass instead of the with-replacement draw that leaves part of a
|
||||
small dataset unseen at low step counts (num_epochs converts to a step budget, but the
|
||||
per-batch index draw is what decides coverage). When a cycle is exhausted the order is
|
||||
reshuffled from the run's own ``rng`` so the index stream stays seed-deterministic and
|
||||
each cycle differs.
|
||||
|
||||
Both trainers share this so the SDXL ``_next_batch`` path and the DiT per-sample draw
|
||||
select indices the same way. Only the index selection changes (with-replacement ->
|
||||
permutation cycles); step count and batch shapes are unchanged.
|
||||
"""
|
||||
|
||||
def __init__(self, n: int, rng: random.Random) -> None:
|
||||
if n <= 0:
|
||||
raise ValueError("PermutationBatchSampler needs at least one item")
|
||||
self._n = n
|
||||
self._rng = rng
|
||||
self._order: list[int] = []
|
||||
self._pos = 0
|
||||
|
||||
def _reshuffle(self) -> None:
|
||||
self._order = list(range(self._n))
|
||||
self._rng.shuffle(self._order)
|
||||
self._pos = 0
|
||||
|
||||
def next_batch(self, k: int) -> list[int]:
|
||||
# k may exceed n (batch larger than the dataset): the permutation is refilled across
|
||||
# as many cycles as needed so the caller always gets exactly k indices and the batch
|
||||
# never shrinks, matching the old sampler's fixed batch shape.
|
||||
out: list[int] = []
|
||||
while len(out) < k:
|
||||
if self._pos >= len(self._order):
|
||||
self._reshuffle()
|
||||
take = min(k - len(out), len(self._order) - self._pos)
|
||||
out.extend(self._order[self._pos : self._pos + take])
|
||||
self._pos += take
|
||||
return out
|
||||
|
||||
|
||||
def discover_image_caption_pairs(
|
||||
data_dir: str | os.PathLike[str],
|
||||
*,
|
||||
|
|
@ -783,6 +843,10 @@ def _write_lora_sidecar(sidecar_path: Path, cfg: DiffusionLoraConfig) -> None:
|
|||
_CONFIG_ALIASES = {
|
||||
"model_name": "base_model",
|
||||
"max_steps": "train_steps",
|
||||
# The generic payload's num_epochs already matches the diffusion field name, but list it
|
||||
# so the epochs override is threaded through the shared-payload path as explicitly as
|
||||
# max_steps -> train_steps is.
|
||||
"num_epochs": "num_epochs",
|
||||
"batch_size": "train_batch_size",
|
||||
"lora_r": "lora_rank",
|
||||
"lr_scheduler_type": "lr_scheduler",
|
||||
|
|
@ -821,6 +885,17 @@ def _config_from_dict(config: dict) -> DiffusionLoraConfig:
|
|||
for k, v in config.items():
|
||||
if k in valid:
|
||||
kwargs[k] = v
|
||||
# Epoch-mode payloads from the generic Studio UI carry max_steps: 0 as the "use epochs"
|
||||
# sentinel, which the max_steps -> train_steps alias copies as train_steps: 0. Since
|
||||
# normalized() rejects train_steps < 1 before resolve_train_steps() can apply num_epochs,
|
||||
# drop a falsy/0 train_steps when num_epochs > 0 so the dataclass default stands in until
|
||||
# epoch resolution replaces it.
|
||||
try:
|
||||
_num_epochs = int(kwargs.get("num_epochs") or 0)
|
||||
except (TypeError, ValueError):
|
||||
_num_epochs = 0
|
||||
if _num_epochs > 0 and not kwargs.get("train_steps"):
|
||||
kwargs.pop("train_steps", None)
|
||||
if kwargs.get("lora_target_modules"):
|
||||
kwargs["lora_target_modules"] = tuple(kwargs["lora_target_modules"])
|
||||
if "gradient_checkpointing" in kwargs:
|
||||
|
|
|
|||
|
|
@ -17,11 +17,14 @@ runs a scripted target on a thread.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import multiprocessing as mp
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Optional
|
||||
|
||||
# Spawn (not fork): a fresh interpreter, matching the LLM training worker, so CUDA/torch
|
||||
|
|
@ -72,6 +75,60 @@ def _default_target(*, event_queue: Any, stop_queue: Any, config: dict) -> None:
|
|||
_METRIC_CAP = 4000
|
||||
|
||||
|
||||
# ── persisted run history ──────────────────────────────────────────────────────
|
||||
# Every terminal run (completed / stopped / error) is recorded as one JSON file --
|
||||
# summary + scrubbed config + the full bounded metric logs -- so the Train tab can show
|
||||
# previous runs like the LLM trainer's history. JSON files (not the LLM sqlite tables)
|
||||
# keep diffusion runs out of the LLM Runs page, whose resume/inspect actions assume an
|
||||
# LLM-shaped run.
|
||||
def _runs_dir() -> Path:
|
||||
from utils.paths.storage_roots import studio_root
|
||||
|
||||
d = studio_root() / "runs" / "diffusion"
|
||||
d.mkdir(parents = True, exist_ok = True)
|
||||
return d
|
||||
|
||||
|
||||
def list_diffusion_runs(limit: int = 20) -> list[dict]:
|
||||
"""Summaries of persisted diffusion runs, newest first. The heavy per-run payload
|
||||
(metric logs, config) stays in the file; fetch it via ``get_diffusion_run``."""
|
||||
try:
|
||||
files = sorted(_runs_dir().glob("*.json"), key = lambda p: p.stat().st_mtime, reverse = True)
|
||||
except Exception: # noqa: BLE001 -- unreadable dir -> no history
|
||||
return []
|
||||
out: list[dict] = []
|
||||
for p in files[: max(0, int(limit))]:
|
||||
try:
|
||||
rec = json.loads(p.read_text(encoding = "utf-8"))
|
||||
except Exception: # noqa: BLE001 -- a corrupt record never breaks the listing
|
||||
continue
|
||||
# A valid-JSON file with the wrong shape (an old or hand-edited record that is not a
|
||||
# dict, or is missing the required string job_id / status) would later blow up the
|
||||
# route's DiffusionTrainingRunSummary(**r); skip it here so one bad record can never
|
||||
# take down the whole Previous runs panel.
|
||||
if not isinstance(rec, dict):
|
||||
continue
|
||||
if not (isinstance(rec.get("job_id"), str) and isinstance(rec.get("status"), str)):
|
||||
continue
|
||||
rec.pop("metric_history", None)
|
||||
rec.pop("config", None)
|
||||
out.append(rec)
|
||||
return out
|
||||
|
||||
|
||||
def get_diffusion_run(job_id: str) -> Optional[dict]:
|
||||
"""The full persisted record for one run (summary + config + metric logs)."""
|
||||
# Records are keyed by the uuid4 hex job id; reject anything else so a crafted id
|
||||
# can never traverse out of the runs directory.
|
||||
if not re.fullmatch(r"[0-9a-f]{32}", str(job_id or "")):
|
||||
return None
|
||||
p = _runs_dir() / f"{job_id}.json"
|
||||
try:
|
||||
return json.loads(p.read_text(encoding = "utf-8"))
|
||||
except Exception: # noqa: BLE001 -- missing/corrupt record
|
||||
return None
|
||||
|
||||
|
||||
def _idle_state() -> dict[str, Any]:
|
||||
return {
|
||||
"active": False,
|
||||
|
|
@ -116,8 +173,8 @@ def _append_metric(
|
|||
Only records finite, positive-step points (mirrors the LLM trainer, which logs history
|
||||
only for step > 0 with a real loss). When the arrays hit ``_METRIC_CAP`` they are
|
||||
decimated in place (keep every other point) so appends stay bounded without losing the
|
||||
curve's shape. lr / grad_norm may be None (kept as None so those series can be
|
||||
sparse while staying index-aligned with ``steps``)."""
|
||||
curve's shape. lr / grad_norm may be None (kept as None so those series can be sparse
|
||||
while staying index-aligned with ``steps``)."""
|
||||
try:
|
||||
istep = int(step)
|
||||
except (TypeError, ValueError):
|
||||
|
|
@ -127,8 +184,9 @@ def _append_metric(
|
|||
floss = _finite_or_none(loss)
|
||||
if floss is None: # non-numeric or non-finite (NaN/Inf): skip, keep the curve JSON-safe
|
||||
return
|
||||
# lr / grad_norm may be None (sparse series) or non-finite; a non-finite value is
|
||||
# nulled, not dropped, so a bad point never taints the (loss-driven) history.
|
||||
# lr / grad_norm may be None (sparse series) or non-finite; non-finite values are
|
||||
# nulled, not dropped, so a bad point never taints the (loss-driven) history while the
|
||||
# arrays stay index-aligned with steps.
|
||||
flr = _finite_or_none(lr)
|
||||
fgn = _finite_or_none(grad_norm)
|
||||
steps = state["metric_steps"]
|
||||
|
|
@ -168,6 +226,8 @@ class DiffusionTrainingService:
|
|||
self._stop_queue: Any = None
|
||||
self._pump: Optional[threading.Thread] = None
|
||||
self._state: dict[str, Any] = _idle_state()
|
||||
# The active job's start config, scrubbed of secrets, kept for the run record.
|
||||
self._config: dict[str, Any] = {}
|
||||
|
||||
# ── lifecycle ────────────────────────────────────────────────────────────
|
||||
def is_active(self) -> bool:
|
||||
|
|
@ -230,6 +290,8 @@ class DiffusionTrainingService:
|
|||
started_at = now,
|
||||
updated_at = now,
|
||||
)
|
||||
# Keep the config (minus secrets) for the persisted run record.
|
||||
self._config = {k: v for k, v in dict(config).items() if k != "hf_token"}
|
||||
self._pump = threading.Thread(
|
||||
target = self._pump_loop, args = (event_queue, self._proc), daemon = True
|
||||
)
|
||||
|
|
@ -291,12 +353,60 @@ class DiffusionTrainingService:
|
|||
updated_at = time.time(),
|
||||
)
|
||||
_ = drained
|
||||
self._persist_run_record()
|
||||
return
|
||||
continue
|
||||
self._apply_event(ev, proc = proc)
|
||||
if ev.get("type") in _TERMINAL:
|
||||
self._persist_run_record()
|
||||
return
|
||||
|
||||
def _persist_run_record(self) -> None:
|
||||
"""Best-effort JSON record of the finished run (summary + scrubbed config + the
|
||||
bounded metric logs) into the studio runs directory. Never fatal: history is a
|
||||
convenience, not part of the training contract."""
|
||||
try:
|
||||
with self._lock:
|
||||
s = dict(self._state)
|
||||
cfg = dict(self._config)
|
||||
if not s.get("job_id") or s.get("status") not in ("completed", "stopped", "error"):
|
||||
return
|
||||
adapter = s.get("output_dir") or cfg.get("output_dir")
|
||||
record = {
|
||||
"job_id": s.get("job_id"),
|
||||
"status": s.get("status"),
|
||||
"message": s.get("message") or "",
|
||||
"family": s.get("family") or cfg.get("model_family"),
|
||||
"base_model": s.get("base_model") or cfg.get("base_model"),
|
||||
"adapter": Path(str(adapter)).name if adapter else None,
|
||||
"instance_prompt": cfg.get("instance_prompt"),
|
||||
"step": s.get("step") or 0,
|
||||
"total_steps": s.get("total_steps") or 0,
|
||||
"loss": s.get("loss"),
|
||||
"avg_loss": s.get("avg_loss"),
|
||||
"learning_rate": s.get("learning_rate"),
|
||||
"grad_norm": s.get("grad_norm"),
|
||||
"samples_per_second": s.get("samples_per_second"),
|
||||
"peak_memory_gb": s.get("peak_memory_gb"),
|
||||
"num_images": s.get("num_images"),
|
||||
"started_at": s.get("started_at"),
|
||||
"ended_at": s.get("updated_at"),
|
||||
"lora_path": s.get("lora_path"),
|
||||
"catalog_path": s.get("catalog_path"),
|
||||
"saved": bool(s.get("lora_path")),
|
||||
"config": cfg,
|
||||
"metric_history": {
|
||||
"steps": s.get("metric_steps") or [],
|
||||
"loss": s.get("metric_loss") or [],
|
||||
"lr": s.get("metric_lr") or [],
|
||||
"grad_norm": s.get("metric_grad_norm") or [],
|
||||
},
|
||||
}
|
||||
path = _runs_dir() / f"{s['job_id']}.json"
|
||||
path.write_text(json.dumps(record), encoding = "utf-8")
|
||||
except Exception: # noqa: BLE001 -- persisting history must never break the run
|
||||
pass
|
||||
|
||||
def _apply_event(
|
||||
self,
|
||||
ev: dict[str, Any],
|
||||
|
|
|
|||
|
|
@ -691,6 +691,15 @@ class DiffusionTrainingStartRequest(BaseModel):
|
|||
1024, ge = 64, le = 2048, description = "Square training resolution (multiple of 8)"
|
||||
)
|
||||
train_steps: int = Field(500, ge = 1, le = 100000)
|
||||
num_epochs: int = Field(
|
||||
0,
|
||||
ge = 0,
|
||||
le = 1000,
|
||||
description = (
|
||||
"0 = use train_steps; > 0 overrides train_steps with epochs x "
|
||||
"ceil(N / (batch x grad_accum)) optimizer steps over the N-image dataset"
|
||||
),
|
||||
)
|
||||
learning_rate: float = Field(1e-4, gt = 0)
|
||||
train_batch_size: int = Field(1, ge = 1, le = 64)
|
||||
gradient_accumulation_steps: int = Field(1, ge = 1, le = 256)
|
||||
|
|
@ -779,8 +788,8 @@ class DiffusionTrainingStatusResponse(BaseModel):
|
|||
loss: Optional[float] = None
|
||||
avg_loss: Optional[float] = None
|
||||
learning_rate: Optional[float] = None
|
||||
# Pre-clip gradient norm from the trainer's progress events (None when clipping is
|
||||
# disabled), feeding the grad-norm chart.
|
||||
# Total pre-clip gradient norm from the last optimizer step (the training health
|
||||
# signal the UI charts alongside the loss).
|
||||
grad_norm: Optional[float] = None
|
||||
num_images: Optional[int] = None
|
||||
in_model_load: bool = False
|
||||
|
|
@ -800,6 +809,43 @@ class DiffusionTrainingStatusResponse(BaseModel):
|
|||
metric_history: Optional[DiffusionMetricHistory] = None
|
||||
|
||||
|
||||
class DiffusionTrainingRunSummary(BaseModel):
|
||||
"""One persisted diffusion training run (terminal), as listed in the Train tab's
|
||||
previous-runs history. The heavy payload (config + metric logs) lives in the detail."""
|
||||
|
||||
job_id: str
|
||||
status: str
|
||||
message: str = ""
|
||||
adapter: Optional[str] = None
|
||||
family: Optional[str] = None
|
||||
base_model: Optional[str] = None
|
||||
step: int = 0
|
||||
total_steps: int = 0
|
||||
avg_loss: Optional[float] = None
|
||||
# Whether this run left an adapter on disk (full completion or stop-and-save).
|
||||
saved: bool = False
|
||||
catalog_path: Optional[str] = None
|
||||
instance_prompt: Optional[str] = None
|
||||
started_at: Optional[float] = None
|
||||
ended_at: Optional[float] = None
|
||||
|
||||
|
||||
class DiffusionTrainingRunDetail(DiffusionTrainingRunSummary):
|
||||
"""The full persisted record: summary + scrubbed start config + metric logs."""
|
||||
|
||||
loss: Optional[float] = None
|
||||
samples_per_second: Optional[float] = None
|
||||
peak_memory_gb: Optional[float] = None
|
||||
num_images: Optional[int] = None
|
||||
lora_path: Optional[str] = None
|
||||
config: Optional[dict] = None
|
||||
metric_history: Optional[DiffusionMetricHistory] = None
|
||||
|
||||
|
||||
class DiffusionTrainingRunsResponse(BaseModel):
|
||||
runs: List[DiffusionTrainingRunSummary] = Field(default_factory = list)
|
||||
|
||||
|
||||
class DiffusionDatasetSummary(BaseModel):
|
||||
"""One image-dataset folder under the Studio datasets root."""
|
||||
|
||||
|
|
|
|||
|
|
@ -70,6 +70,9 @@ from models.training import (
|
|||
DiffusionMetricHistory,
|
||||
DiffusionTrainableFamily,
|
||||
DiffusionTrainingInfoResponse,
|
||||
DiffusionTrainingRunDetail,
|
||||
DiffusionTrainingRunsResponse,
|
||||
DiffusionTrainingRunSummary,
|
||||
DiffusionTrainingStartRequest,
|
||||
DiffusionTrainingStartResponse,
|
||||
DiffusionTrainingStatusResponse,
|
||||
|
|
@ -77,6 +80,7 @@ from models.training import (
|
|||
)
|
||||
from models.responses import TrainingStopResponse, TrainingMetricsResponse
|
||||
from pydantic import BaseModel as PydanticBaseModel
|
||||
from pydantic import ValidationError
|
||||
|
||||
|
||||
class TrainingStopRequest(PydanticBaseModel):
|
||||
|
|
@ -1335,6 +1339,49 @@ async def diffusion_training_status(current_subject: str = Depends(get_current_s
|
|||
return DiffusionTrainingStatusResponse(**snap, metric_history = metric_history)
|
||||
|
||||
|
||||
@router.get("/diffusion/runs", response_model = DiffusionTrainingRunsResponse)
|
||||
async def list_diffusion_training_runs(
|
||||
limit: int = 20, current_subject: str = Depends(get_current_subject)
|
||||
):
|
||||
"""Previous diffusion training runs (terminal), newest first, from the persisted
|
||||
per-run records. Summaries only; fetch one run for its config + metric logs."""
|
||||
from core.training.diffusion_training_service import list_diffusion_runs
|
||||
|
||||
summaries: list[DiffusionTrainingRunSummary] = []
|
||||
for r in list_diffusion_runs(limit = limit):
|
||||
# list_diffusion_runs already skips non-dict / missing-id records, but a record with
|
||||
# a wrong-typed field (e.g. a non-numeric avg_loss) would still raise here; catch it
|
||||
# per record so one bad file never breaks the whole Previous runs panel.
|
||||
try:
|
||||
summaries.append(DiffusionTrainingRunSummary(**r))
|
||||
except ValidationError:
|
||||
continue
|
||||
return DiffusionTrainingRunsResponse(runs = summaries)
|
||||
|
||||
|
||||
@router.get("/diffusion/runs/{job_id}", response_model = DiffusionTrainingRunDetail)
|
||||
async def get_diffusion_training_run(
|
||||
job_id: str, current_subject: str = Depends(get_current_subject)
|
||||
):
|
||||
"""One persisted diffusion run's full record: summary + scrubbed start config + the
|
||||
step/loss/grad-norm logs (for re-plotting a past run's charts)."""
|
||||
from core.training.diffusion_training_service import get_diffusion_run
|
||||
|
||||
rec = get_diffusion_run(job_id)
|
||||
# A valid-JSON file that is not an object (a truncated / hand-edited [] record) would make
|
||||
# DiffusionTrainingRunDetail(**rec) raise TypeError -- not the ValidationError caught below
|
||||
# -- and 500 the endpoint. Treat any non-dict record as absent, matching the list route's
|
||||
# shape check.
|
||||
if not isinstance(rec, dict):
|
||||
raise HTTPException(status_code = 404, detail = "No such training run.")
|
||||
try:
|
||||
return DiffusionTrainingRunDetail(**rec)
|
||||
except ValidationError:
|
||||
# A malformed on-disk record (hand-edited / older shape) should read as absent
|
||||
# rather than 500 the endpoint, mirroring how the list route skips bad records.
|
||||
raise HTTPException(status_code = 404, detail = "No such training run.")
|
||||
|
||||
|
||||
# Extensions accepted into an image-training dataset folder: images the trainer reads,
|
||||
# plus its caption sources (per-image sidecars and metadata/captions jsonl).
|
||||
_DIFFUSION_DATASET_IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ from core.training.diffusion_lora_trainer import (
|
|||
_config_from_dict,
|
||||
compute_sdxl_add_time_ids,
|
||||
discover_image_caption_pairs,
|
||||
resolve_train_steps,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -114,6 +115,60 @@ def test_config_normalized_validation(kw):
|
|||
DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", **kw).normalized()
|
||||
|
||||
|
||||
def _cfg(**kw):
|
||||
return DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", **kw)
|
||||
|
||||
|
||||
def test_resolve_train_steps_uses_train_steps_when_epochs_disabled():
|
||||
# num_epochs == 0 leaves the explicit train_steps untouched, whatever the image count.
|
||||
cfg = _cfg(train_steps = 300, num_epochs = 0)
|
||||
assert resolve_train_steps(cfg, 20) == 300
|
||||
assert resolve_train_steps(cfg, 1) == 300
|
||||
|
||||
|
||||
def test_resolve_train_steps_epochs_ceil_over_batch_and_grad_accum():
|
||||
# One epoch = ceil(N / (batch x grad_accum)) optimizer steps; num_epochs multiplies it.
|
||||
# 10 images, batch 4, grad_accum 1 -> ceil(10/4)=3 steps/epoch.
|
||||
assert resolve_train_steps(_cfg(num_epochs = 1, train_batch_size = 4), 10) == 3
|
||||
assert resolve_train_steps(_cfg(num_epochs = 5, train_batch_size = 4), 10) == 15
|
||||
# grad_accum widens the effective batch: 100 images, batch 2, grad_accum 3 -> per_step=6,
|
||||
# ceil(100/6)=17 steps/epoch, 2 epochs -> 34.
|
||||
cfg = _cfg(num_epochs = 2, train_batch_size = 2, gradient_accumulation_steps = 3)
|
||||
assert resolve_train_steps(cfg, 100) == 34
|
||||
# An exact multiple does not round up: 8 images / batch 4 -> 2 steps/epoch.
|
||||
assert resolve_train_steps(_cfg(num_epochs = 3, train_batch_size = 4), 8) == 6
|
||||
|
||||
|
||||
def test_resolve_train_steps_single_image_dataset():
|
||||
# A one-image dataset is one optimizer step per epoch, so num_epochs == steps.
|
||||
assert resolve_train_steps(_cfg(num_epochs = 7, train_batch_size = 4), 1) == 7
|
||||
|
||||
|
||||
def test_resolve_train_steps_caps_at_100000():
|
||||
# The run length is capped at 100000 even for absurd epoch counts (matches the request
|
||||
# model's train_steps ceiling), so a huge epochs x dataset never overflows the loop.
|
||||
cfg = _cfg(num_epochs = 1000, train_batch_size = 1)
|
||||
assert resolve_train_steps(cfg, 10_000) == 100000
|
||||
|
||||
|
||||
def test_config_normalized_num_epochs_bounds():
|
||||
# 0 (disabled) and the 1..1000 range normalise; out-of-range is rejected.
|
||||
assert _cfg(num_epochs = 0).normalized().num_epochs == 0
|
||||
assert _cfg(num_epochs = 1000).normalized().num_epochs == 1000
|
||||
with pytest.raises(ValueError, match = "num_epochs"):
|
||||
_cfg(num_epochs = -1).normalized()
|
||||
with pytest.raises(ValueError, match = "num_epochs"):
|
||||
_cfg(num_epochs = 1001).normalized()
|
||||
|
||||
|
||||
def test_config_from_dict_threads_num_epochs():
|
||||
# num_epochs flows through the shared-payload adapter onto the diffusion field.
|
||||
cfg = _config_from_dict(
|
||||
{"base_model": "b", "data_dir": "d", "output_dir": "o", "num_epochs": 12}
|
||||
)
|
||||
assert cfg.num_epochs == 12
|
||||
|
||||
|
||||
def test_compute_sdxl_add_time_ids():
|
||||
assert compute_sdxl_add_time_ids(1024) == (1024, 1024, 0, 0, 1024, 1024)
|
||||
|
||||
|
|
|
|||
|
|
@ -66,6 +66,18 @@ class _FakeCtx:
|
|||
return _FakeProc(target, kwargs, daemon)
|
||||
|
||||
|
||||
@pytest.fixture(autouse = True)
|
||||
def _isolated_runs_dir(monkeypatch, tmp_path):
|
||||
"""Terminal service events persist a run record; point the runs dir at tmp so tests
|
||||
never write into a real studio home. Yields the dir for the history tests."""
|
||||
import core.training.diffusion_training_service as dts
|
||||
|
||||
d = tmp_path / "runs" / "diffusion"
|
||||
d.mkdir(parents = True, exist_ok = True)
|
||||
monkeypatch.setattr(dts, "_runs_dir", lambda: d)
|
||||
yield d
|
||||
|
||||
|
||||
def _happy_target(*, event_queue, stop_queue, config):
|
||||
event_queue.put({"type": "model_load_started", "num_images": 3})
|
||||
event_queue.put({"type": "model_load_completed"})
|
||||
|
|
@ -206,12 +218,15 @@ def test_progress_nulls_non_finite_floats_for_strict_json():
|
|||
"loss": float("nan"),
|
||||
"avg_loss": float("inf"),
|
||||
"learning_rate": float("-inf"),
|
||||
"grad_norm": float("inf"),
|
||||
}
|
||||
)
|
||||
snap = svc.status()
|
||||
assert snap["loss"] is None
|
||||
assert snap["avg_loss"] is None
|
||||
assert snap["learning_rate"] is None
|
||||
# The reviewer's exact case: an inf pre-clip grad norm must not reach the status JSON.
|
||||
assert snap["grad_norm"] is None
|
||||
# The non-finite point is skipped in the history, so the loss series stays clean.
|
||||
assert snap["metric_loss"] == []
|
||||
assert snap["metric_steps"] == []
|
||||
|
|
@ -355,6 +370,106 @@ def test_route_start_forwards_extra_training_knobs(client):
|
|||
assert client._fake.started_with["lora_target_modules"] == ["to_q", "to_v"]
|
||||
|
||||
|
||||
def test_route_start_forwards_num_epochs(client):
|
||||
# Epochs mode: the frontend omits train_steps and sends num_epochs; it must reach the
|
||||
# service so the trainer can resolve it against the dataset size.
|
||||
body = {k: v for k, v in _BODY.items() if k != "train_steps"}
|
||||
r = client.post("/api/train/diffusion/start", json = {**body, "num_epochs": 8})
|
||||
assert r.status_code == 200, r.text
|
||||
assert client._fake.started_with["num_epochs"] == 8
|
||||
|
||||
|
||||
def test_request_model_num_epochs_bounds():
|
||||
# The request schema mirrors DiffusionLoraConfig's 0..1000 num_epochs range.
|
||||
from pydantic import ValidationError
|
||||
|
||||
from models.training import DiffusionTrainingStartRequest
|
||||
|
||||
base = {"base_model": "b", "data_dir": "d", "output_dir": "o"}
|
||||
assert DiffusionTrainingStartRequest(**base).num_epochs == 0 # default = use train_steps
|
||||
assert DiffusionTrainingStartRequest(**base, num_epochs = 1000).num_epochs == 1000
|
||||
for bad in (-1, 1001):
|
||||
with pytest.raises(ValidationError):
|
||||
DiffusionTrainingStartRequest(**base, num_epochs = bad)
|
||||
|
||||
|
||||
def test_config_from_dict_epoch_mode_drops_max_steps_sentinel():
|
||||
# The generic Studio epoch-mode payload sends max_steps: 0 as the "use epochs" sentinel.
|
||||
# The max_steps -> train_steps alias would copy that 0 and normalized() would reject
|
||||
# train_steps < 1 before epochs are resolved; _config_from_dict must drop the falsy
|
||||
# value so the default train_steps stands in until resolve_train_steps applies num_epochs.
|
||||
from core.training.diffusion_train_common import DiffusionLoraConfig, _config_from_dict
|
||||
|
||||
cfg = _config_from_dict(
|
||||
{
|
||||
"base_model": "stabilityai/stable-diffusion-xl-base-1.0",
|
||||
"data_dir": "d",
|
||||
"output_dir": "o",
|
||||
"max_steps": 0,
|
||||
"num_epochs": 2,
|
||||
}
|
||||
)
|
||||
# 0 was dropped: the dataclass default train_steps stands in and num_epochs carries over.
|
||||
assert cfg.train_steps == DiffusionLoraConfig.train_steps
|
||||
assert cfg.num_epochs == 2
|
||||
# normalized() no longer raises on the epoch-mode payload.
|
||||
norm = cfg.normalized()
|
||||
assert norm.num_epochs == 2
|
||||
|
||||
# An explicit non-zero max_steps in epochs mode is still honored (only the 0 sentinel is
|
||||
# dropped), and a plain steps payload (no num_epochs) keeps max_steps: 0 -> train_steps 0
|
||||
# so normalized() surfaces the invalid value as before.
|
||||
cfg_explicit = _config_from_dict(
|
||||
{
|
||||
"base_model": "stabilityai/stable-diffusion-xl-base-1.0",
|
||||
"data_dir": "d",
|
||||
"output_dir": "o",
|
||||
"max_steps": 25,
|
||||
"num_epochs": 2,
|
||||
}
|
||||
)
|
||||
assert cfg_explicit.train_steps == 25
|
||||
|
||||
|
||||
def test_permutation_sampler_covers_dataset_once_per_cycle():
|
||||
# Every index must appear exactly once per cycle before any repeat (epoch-style pass),
|
||||
# so a short run over a small dataset never leaves images unseen the way the old
|
||||
# with-replacement draw did. Consecutive cycles must be reshuffled (differ).
|
||||
import random
|
||||
|
||||
from core.training.diffusion_train_common import PermutationBatchSampler
|
||||
|
||||
n = 100
|
||||
sampler = PermutationBatchSampler(n, random.Random(0))
|
||||
|
||||
# Draw exactly one cycle in batches of 3 (n not divisible by the batch, so a batch spans
|
||||
# the cycle boundary); the first n indices must be a permutation of range(n).
|
||||
drawn: list[int] = []
|
||||
while len(drawn) < n:
|
||||
drawn.extend(sampler.next_batch(3))
|
||||
first_cycle = drawn[:n]
|
||||
assert sorted(first_cycle) == list(range(n)) # each index once, none missing
|
||||
|
||||
# The next full cycle is also a permutation, and it is reshuffled (order differs).
|
||||
fresh = PermutationBatchSampler(n, random.Random(0))
|
||||
cycle_a = fresh.next_batch(n)
|
||||
cycle_b = fresh.next_batch(n)
|
||||
assert sorted(cycle_a) == list(range(n))
|
||||
assert sorted(cycle_b) == list(range(n))
|
||||
assert cycle_a != cycle_b # cycles are reshuffled, not repeated in the same order
|
||||
|
||||
# A seed replays the exact index stream (determinism for reproducible runs).
|
||||
replay = PermutationBatchSampler(n, random.Random(0))
|
||||
assert replay.next_batch(n) == cycle_a
|
||||
|
||||
# A batch larger than the dataset refills across cycles so it never shrinks (batch shape
|
||||
# preserved), even though it must then repeat indices within the batch.
|
||||
big = PermutationBatchSampler(4, random.Random(1))
|
||||
batch = big.next_batch(10)
|
||||
assert len(batch) == 10
|
||||
assert set(batch) == {0, 1, 2, 3}
|
||||
|
||||
|
||||
def test_route_start_accepts_zero_max_grad_norm(client):
|
||||
# 0 is the documented "disable clipping" value (the trainer skips clip_grad_norm_);
|
||||
# the request model must not reject it.
|
||||
|
|
@ -780,3 +895,164 @@ def test_start_ungated_base_preflight_is_noop(client, monkeypatch):
|
|||
)
|
||||
assert r.status_code == 200, r.text
|
||||
assert client._fake.started_with["base_model"] == "black-forest-labs/FLUX.1-dev"
|
||||
|
||||
|
||||
# ── persisted run history ──────────────────────────────────────────────────────
|
||||
def test_run_record_persisted_on_complete(_isolated_runs_dir):
|
||||
# A completed run writes one JSON record: summary + scrubbed config + metric logs.
|
||||
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
|
||||
job_id = svc.start({**_CFG, "model_family": "z-image", "hf_token": "SECRET"})
|
||||
_wait_status(svc, "completed")
|
||||
# The pump persists right after the terminal event; give the thread a beat.
|
||||
time.sleep(0.1)
|
||||
|
||||
import json
|
||||
|
||||
rec = json.loads((_isolated_runs_dir / f"{job_id}.json").read_text())
|
||||
assert rec["job_id"] == job_id
|
||||
assert rec["status"] == "completed"
|
||||
assert rec["saved"] is True
|
||||
assert rec["adapter"] == "out" # basename of /tmp/out
|
||||
assert rec["family"] == "z-image" # falls back to the config's model_family
|
||||
assert rec["step"] == 2 and rec["total_steps"] == 2
|
||||
assert rec["avg_loss"] == 0.45
|
||||
assert rec["metric_history"]["steps"] == [1, 2]
|
||||
assert rec["metric_history"]["loss"] == [0.5, 0.4]
|
||||
# Secrets never land on disk.
|
||||
assert "hf_token" not in rec["config"]
|
||||
assert rec["config"]["model_family"] == "z-image"
|
||||
|
||||
|
||||
def test_run_record_no_save_stop_marks_unsaved(_isolated_runs_dir):
|
||||
# A cancel (stop without save) persists too, flagged as not saved.
|
||||
def _cancel_target(*, event_queue, stop_queue, config):
|
||||
event_queue.put({"type": "model_load_completed"})
|
||||
stop_queue.get(timeout = 5.0)
|
||||
event_queue.put(
|
||||
{"type": "complete", "output_dir": None, "lora_path": None, "stopped": True}
|
||||
)
|
||||
|
||||
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _cancel_target)
|
||||
job_id = svc.start(dict(_CFG))
|
||||
_wait_status(svc, "running")
|
||||
svc.stop(save = False)
|
||||
_wait_status(svc, "stopped")
|
||||
time.sleep(0.1)
|
||||
|
||||
import json
|
||||
|
||||
rec = json.loads((_isolated_runs_dir / f"{job_id}.json").read_text())
|
||||
assert rec["status"] == "stopped"
|
||||
assert rec["saved"] is False and rec["lora_path"] is None
|
||||
|
||||
|
||||
def test_runs_endpoints_list_and_detail(client, _isolated_runs_dir):
|
||||
# Seed two records directly (the endpoints read the persisted files, not the service).
|
||||
import json
|
||||
import os
|
||||
|
||||
a = {
|
||||
"job_id": "a" * 32,
|
||||
"status": "completed",
|
||||
"adapter": "first",
|
||||
"saved": True,
|
||||
"step": 10,
|
||||
"total_steps": 10,
|
||||
"avg_loss": 0.4,
|
||||
"config": {"train_steps": 10},
|
||||
"metric_history": {"steps": [1], "loss": [0.4], "lr": [1e-4], "grad_norm": [0.2]},
|
||||
}
|
||||
b = {
|
||||
"job_id": "b" * 32,
|
||||
"status": "stopped",
|
||||
"adapter": "second",
|
||||
"saved": False,
|
||||
"step": 3,
|
||||
"total_steps": 10,
|
||||
"avg_loss": 0.6,
|
||||
"config": {"train_steps": 10},
|
||||
"metric_history": {"steps": [1], "loss": [0.6], "lr": [1e-4], "grad_norm": [0.3]},
|
||||
}
|
||||
pa = _isolated_runs_dir / f"{a['job_id']}.json"
|
||||
pb = _isolated_runs_dir / f"{b['job_id']}.json"
|
||||
pa.write_text(json.dumps(a))
|
||||
pb.write_text(json.dumps(b))
|
||||
os.utime(pa, (1000, 1000))
|
||||
os.utime(pb, (2000, 2000)) # b is newer -> listed first
|
||||
|
||||
r = client.get("/api/train/diffusion/runs")
|
||||
assert r.status_code == 200, r.text
|
||||
runs = r.json()["runs"]
|
||||
assert [x["adapter"] for x in runs] == ["second", "first"]
|
||||
# Summaries stay light: no config / metric logs.
|
||||
assert "config" not in runs[0] and "metric_history" not in runs[0]
|
||||
|
||||
r = client.get(f"/api/train/diffusion/runs/{a['job_id']}")
|
||||
assert r.status_code == 200, r.text
|
||||
detail = r.json()
|
||||
assert detail["adapter"] == "first"
|
||||
assert detail["metric_history"]["grad_norm"] == [0.2]
|
||||
assert detail["config"] == {"train_steps": 10}
|
||||
|
||||
# Unknown and malformed ids 404 (malformed also covers path traversal).
|
||||
assert client.get(f"/api/train/diffusion/runs/{'c' * 32}").status_code == 404
|
||||
assert client.get("/api/train/diffusion/runs/not-a-job-id").status_code == 404
|
||||
|
||||
|
||||
def test_list_diffusion_runs_skips_wrong_shape_records(_isolated_runs_dir):
|
||||
# A valid-JSON file with the wrong shape (non-dict, or missing the required string
|
||||
# job_id / status) must be skipped by list_diffusion_runs so it never reaches the route's
|
||||
# DiffusionTrainingRunSummary(**r) and takes down the whole Previous runs panel.
|
||||
import json
|
||||
|
||||
from core.training.diffusion_training_service import list_diffusion_runs
|
||||
|
||||
good = {"job_id": "a" * 32, "status": "completed", "adapter": "good", "saved": True}
|
||||
(_isolated_runs_dir / "good.json").write_text(json.dumps(good))
|
||||
# A JSON list (not a dict).
|
||||
(_isolated_runs_dir / "not_a_dict.json").write_text(json.dumps([1, 2, 3]))
|
||||
# A dict missing the required job_id / status.
|
||||
(_isolated_runs_dir / "no_ids.json").write_text(json.dumps({"adapter": "orphan"}))
|
||||
# A dict whose job_id / status are the wrong type.
|
||||
(_isolated_runs_dir / "bad_types.json").write_text(
|
||||
json.dumps({"job_id": 123, "status": None, "adapter": "typed"})
|
||||
)
|
||||
|
||||
runs = list_diffusion_runs()
|
||||
adapters = [r.get("adapter") for r in runs]
|
||||
assert adapters == ["good"] # only the well-shaped record survives
|
||||
|
||||
|
||||
def test_runs_route_tolerates_bad_field_record(client, _isolated_runs_dir):
|
||||
# A record that passes the service's shape check but has a wrong-typed field (a
|
||||
# non-numeric avg_loss) would raise pydantic ValidationError in the route; the route must
|
||||
# catch it per record so one bad file never breaks the panel and the good runs still list.
|
||||
import json
|
||||
|
||||
good = {"job_id": "a" * 32, "status": "completed", "adapter": "good", "saved": True}
|
||||
bad = {
|
||||
"job_id": "b" * 32,
|
||||
"status": "completed",
|
||||
"adapter": "bad",
|
||||
"avg_loss": "not-a-number", # str where the summary expects Optional[float]
|
||||
}
|
||||
(_isolated_runs_dir / f"{good['job_id']}.json").write_text(json.dumps(good))
|
||||
(_isolated_runs_dir / f"{bad['job_id']}.json").write_text(json.dumps(bad))
|
||||
|
||||
r = client.get("/api/train/diffusion/runs")
|
||||
assert r.status_code == 200, r.text
|
||||
adapters = [x["adapter"] for x in r.json()["runs"]]
|
||||
assert adapters == ["good"] # the bad-field record was skipped, the good one remained
|
||||
|
||||
|
||||
def test_run_detail_route_non_object_record_is_404(client, _isolated_runs_dir):
|
||||
# A valid-JSON but non-object record (a truncated / hand-edited [] file named with a real
|
||||
# job id) makes DiffusionTrainingRunDetail(**rec) raise TypeError, not ValidationError; the
|
||||
# detail route must shape-check like the list path and 404 instead of 500.
|
||||
import json
|
||||
|
||||
job_id = "a" * 32
|
||||
(_isolated_runs_dir / f"{job_id}.json").write_text(json.dumps([]))
|
||||
|
||||
r = client.get(f"/api/train/diffusion/runs/{job_id}")
|
||||
assert r.status_code == 404, r.text
|
||||
|
|
|
|||
|
|
@ -270,6 +270,9 @@ export interface DiffusionTrainingStartRequest {
|
|||
instance_prompt?: string | null;
|
||||
resolution?: number;
|
||||
train_steps?: number;
|
||||
// 0 or omitted uses train_steps. > 0 overrides train_steps with that many epochs
|
||||
// (full passes over the dataset, in optimizer steps).
|
||||
num_epochs?: number;
|
||||
learning_rate?: number;
|
||||
train_batch_size?: number;
|
||||
gradient_accumulation_steps?: number;
|
||||
|
|
@ -281,6 +284,19 @@ export interface DiffusionTrainingStartRequest {
|
|||
mixed_precision?: "bf16" | "fp16" | "no";
|
||||
gradient_checkpointing?: boolean;
|
||||
lr_scheduler?: string;
|
||||
lr_warmup_steps?: number;
|
||||
// DiT-family quantised base precision (nf4 QLoRA by default). Ignored for sdxl, which
|
||||
// uses mixed_precision instead. "auto" lets the backend pick per family.
|
||||
base_precision?: "nf4" | "bf16" | "int8" | "fp8" | "auto";
|
||||
// Whether to torch.compile the transformer (DiT families that support it). "auto" lets
|
||||
// the backend decide; "off"/"on" force it.
|
||||
compile_transformer?: "off" | "on" | "auto";
|
||||
// Precompute + cache the VAE latents before the loop (skips re-encoding each epoch).
|
||||
cache_latents?: boolean;
|
||||
// How many augmentation variants to cache per image when caching latents (1..16).
|
||||
cache_variants?: number;
|
||||
// Allow TF32 matmuls on Ampere+ for a throughput win at negligible quality cost.
|
||||
enable_tf32?: boolean;
|
||||
// Forwarded to the pipeline's from_pretrained for a gated/private base repo (e.g. FLUX).
|
||||
hf_token?: string | null;
|
||||
}
|
||||
|
|
@ -291,7 +307,8 @@ export interface DiffusionMetricHistory {
|
|||
steps: number[];
|
||||
loss: number[];
|
||||
lr: Array<number | null>;
|
||||
grad_norm: Array<number | null>;
|
||||
// Total pre-clip gradient norm per step (the training health signal the charts show).
|
||||
grad_norm?: Array<number | null>;
|
||||
}
|
||||
|
||||
// A snapshot of the current diffusion training job (GET /api/train/diffusion/status).
|
||||
|
|
@ -305,6 +322,7 @@ export interface DiffusionTrainingStatus {
|
|||
loss: number | null;
|
||||
avg_loss: number | null;
|
||||
learning_rate: number | null;
|
||||
grad_norm?: number | null;
|
||||
num_images: number | null;
|
||||
in_model_load: boolean;
|
||||
output_dir: string | null;
|
||||
|
|
@ -335,8 +353,57 @@ export async function startDiffusionTraining(
|
|||
);
|
||||
}
|
||||
|
||||
export async function stopDiffusionTraining(): Promise<{ status: string }> {
|
||||
return parseJson(await authFetch("/api/train/diffusion/stop", { method: "POST" }));
|
||||
// Request a stop of the running job. `save` (default true) writes the current adapter
|
||||
// before halting ("Stop and save"); false discards it ("Stop").
|
||||
export async function stopDiffusionTraining(save = true): Promise<{ status: string }> {
|
||||
return parseJson(
|
||||
await authFetch("/api/train/diffusion/stop", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ save }),
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
// One persisted (terminal) diffusion training run, as listed in the previous-runs
|
||||
// history. The detail adds the scrubbed start config + the full metric logs.
|
||||
export interface DiffusionTrainingRunSummary {
|
||||
job_id: string;
|
||||
status: string;
|
||||
message?: string;
|
||||
adapter?: string | null;
|
||||
family?: string | null;
|
||||
base_model?: string | null;
|
||||
step: number;
|
||||
total_steps: number;
|
||||
avg_loss?: number | null;
|
||||
saved: boolean;
|
||||
catalog_path?: string | null;
|
||||
instance_prompt?: string | null;
|
||||
started_at?: number | null;
|
||||
ended_at?: number | null;
|
||||
}
|
||||
|
||||
export interface DiffusionTrainingRunDetail extends DiffusionTrainingRunSummary {
|
||||
loss?: number | null;
|
||||
samples_per_second?: number | null;
|
||||
peak_memory_gb?: number | null;
|
||||
num_images?: number | null;
|
||||
lora_path?: string | null;
|
||||
config?: Record<string, unknown> | null;
|
||||
metric_history?: DiffusionMetricHistory | null;
|
||||
}
|
||||
|
||||
export async function listDiffusionTrainingRuns(
|
||||
limit = 20,
|
||||
): Promise<{ runs: DiffusionTrainingRunSummary[] }> {
|
||||
return parseJson(await authFetch(`/api/train/diffusion/runs?limit=${limit}`));
|
||||
}
|
||||
|
||||
export async function getDiffusionTrainingRun(
|
||||
jobId: string,
|
||||
): Promise<DiffusionTrainingRunDetail> {
|
||||
return parseJson(await authFetch(`/api/train/diffusion/runs/${encodeURIComponent(jobId)}`));
|
||||
}
|
||||
|
||||
export async function getDiffusionTrainingStatus(): Promise<DiffusionTrainingStatus> {
|
||||
|
|
@ -369,6 +436,13 @@ export interface DiffusionTrainableFamily {
|
|||
} | null;
|
||||
vram_note?: string | null;
|
||||
gated?: boolean | null;
|
||||
// Quantised base precisions this family can train in (subset of
|
||||
// ["nf4","bf16","int8","fp8","auto"]); empty for sdxl, which uses mixed_precision.
|
||||
precision_modes?: string[];
|
||||
// The precision the backend recommends for this family (marked "(recommended)").
|
||||
recommended_precision?: string;
|
||||
// Whether the family's transformer can be torch.compile'd (gates the Speed > Compile row).
|
||||
supports_compile?: boolean;
|
||||
}
|
||||
|
||||
// Where diffusion training reads/writes on this Studio, plus usable dataset folders.
|
||||
|
|
|
|||
|
|
@ -1955,21 +1955,22 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
|
|||
shared element matches. The load progress shows in a chat-style toast,
|
||||
not here. ── */}
|
||||
<div className="flex h-[48px] shrink-0 items-start justify-between pl-2 pr-2 pt-[11px]">
|
||||
<ModelSelector
|
||||
models={MODELS}
|
||||
value={status?.loaded ? status.repo_id ?? undefined : undefined}
|
||||
activeGgufVariant={quant}
|
||||
onValueChange={handleModelSelect}
|
||||
onEject={status?.loaded ? handleUnload : undefined}
|
||||
variant="ghost"
|
||||
className="!h-[34px]"
|
||||
task={IMAGE_GEN_TASKS}
|
||||
open={active && selectorOpen}
|
||||
onOpenChange={(o) => setSelectorOpen(active && o)}
|
||||
/>
|
||||
<div className="flex items-center gap-2">
|
||||
{/* Create | Train page-mode switch, next to the model selector. Create is the
|
||||
generation workspace; Train is the full-page LoRA training workspace. */}
|
||||
<ModelSelector
|
||||
models={MODELS}
|
||||
value={status?.loaded ? status.repo_id ?? undefined : undefined}
|
||||
activeGgufVariant={quant}
|
||||
onValueChange={handleModelSelect}
|
||||
onEject={status?.loaded ? handleUnload : undefined}
|
||||
variant="ghost"
|
||||
className="!h-[34px]"
|
||||
task={IMAGE_GEN_TASKS}
|
||||
open={active && selectorOpen}
|
||||
onOpenChange={(o) => setSelectorOpen(active && o)}
|
||||
/>
|
||||
{/* Create | Train page-mode switch, on the left next to the model selector
|
||||
(the selector itself stays leftmost: its position is shared with Chat's).
|
||||
Create is the generation workspace; Train is the LoRA training workspace. */}
|
||||
<Tabs value={pageMode} onValueChange={(v) => setPageMode(v as "create" | "train")}>
|
||||
<TabsList className="h-[34px]">
|
||||
<TabsTrigger value="create" className="w-[64px]">
|
||||
|
|
@ -1980,6 +1981,8 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
|
|||
</TabsTrigger>
|
||||
</TabsList>
|
||||
</Tabs>
|
||||
</div>
|
||||
<div className="flex items-center gap-2">
|
||||
{/* Single fixed toggle for the right-docked Advanced panel (mirrors Chat's settings
|
||||
toggle, same icon in both states so it never moves). Highlighted when open.
|
||||
Only meaningful in Create mode (load-time tuning), so hidden while training. */}
|
||||
|
|
|
|||
|
|
@ -4,15 +4,14 @@
|
|||
import { type ReactElement, useMemo } from "react";
|
||||
|
||||
import type { TrainingSeriesPoint } from "@/features/training";
|
||||
// The loss + LR cards are pure presentational (props only), so reuse them directly. We do
|
||||
// NOT reuse ChartsSection/ChartsContent: those also render Grad Norm and an Eval Loss card,
|
||||
// which are meaningless for diffusion LoRA training and showed as an empty card and an
|
||||
// "Evaluation not configured" placeholder. This is a diffusion-only two-card layout.
|
||||
// The loss + grad-norm cards are pure presentational (props only), so reuse them directly.
|
||||
// We do NOT reuse ChartsSection/ChartsContent: those also render an LR and an Eval Loss
|
||||
// card, which add little for diffusion LoRA training (the LR curve is the deterministic
|
||||
// schedule the user just picked; eval is not configured). This is a diffusion-only
|
||||
// two-card layout: Training Loss + Grad Norm (the actual training health signal).
|
||||
// eslint-disable-next-line no-restricted-imports
|
||||
import { GradNormChartCard } from "@/features/studio/sections/charts/grad-norm-chart-card";
|
||||
// eslint-disable-next-line no-restricted-imports
|
||||
import { LearningRateChartCard } from "@/features/studio/sections/charts/learning-rate-chart-card";
|
||||
// eslint-disable-next-line no-restricted-imports
|
||||
import { TrainingLossChartCard } from "@/features/studio/sections/charts/training-loss-chart-card";
|
||||
// eslint-disable-next-line no-restricted-imports
|
||||
import {
|
||||
|
|
@ -45,18 +44,17 @@ function fullStepDomain(steps: number[]): [number, number] {
|
|||
return [min, max];
|
||||
}
|
||||
|
||||
// A diffusion-only metrics view: Training Loss and Learning Rate side by side, plus Grad
|
||||
// Norm (the pre-clip total gradient norm; spikes flag instability that raw loss noise
|
||||
// hides), with a note under the loss card explaining why per-step loss looks noisy.
|
||||
// A diffusion-only metrics view: Training Loss and Grad Norm, side by side, with a note
|
||||
// under the loss card explaining why per-step loss looks noisy. Always renders both cards
|
||||
// (even with no data) so the parent can decide when to mount them; we never early-return
|
||||
// null here.
|
||||
export function DiffusionCharts({
|
||||
lossHistory,
|
||||
lrHistory,
|
||||
gradNormHistory = [],
|
||||
gradNormHistory,
|
||||
}: {
|
||||
lossHistory: TrainingSeriesPoint[];
|
||||
lrHistory: TrainingSeriesPoint[];
|
||||
gradNormHistory?: TrainingSeriesPoint[];
|
||||
}): ReactElement | null {
|
||||
gradNormHistory: TrainingSeriesPoint[];
|
||||
}): ReactElement {
|
||||
const lossItems = useMemo(() => toLossItems(lossHistory), [lossHistory]);
|
||||
const smoothed = useMemo(
|
||||
() => (lossItems.length > 0 ? ema(lossItems, SMOOTHING) : []),
|
||||
|
|
@ -76,18 +74,7 @@ export function DiffusionCharts({
|
|||
[reducedLoss],
|
||||
);
|
||||
|
||||
const lrData = useMemo(
|
||||
() =>
|
||||
compressSeries(
|
||||
lrHistory
|
||||
.filter((p) => Number.isFinite(p.value))
|
||||
.map((p) => ({ step: p.step, lr: p.value, displayLr: p.value })),
|
||||
MAX_RENDER_POINTS,
|
||||
),
|
||||
[lrHistory],
|
||||
);
|
||||
|
||||
const gradNormData = useMemo(
|
||||
const gradData = useMemo(
|
||||
() =>
|
||||
compressSeries(
|
||||
gradNormHistory
|
||||
|
|
@ -101,10 +88,9 @@ export function DiffusionCharts({
|
|||
const steps = useMemo(() => {
|
||||
const set = new Set<number>();
|
||||
for (const p of lossData) set.add(p.step);
|
||||
for (const p of lrData) set.add(p.step);
|
||||
for (const p of gradNormData) set.add(p.step);
|
||||
for (const p of gradData) set.add(p.step);
|
||||
return Array.from(set).sort((a, b) => a - b);
|
||||
}, [lossData, lrData, gradNormData]);
|
||||
}, [lossData, gradData]);
|
||||
|
||||
const stepDomain = useMemo(() => fullStepDomain(steps), [steps]);
|
||||
const xAxisTicks = useMemo(
|
||||
|
|
@ -116,22 +102,15 @@ export function DiffusionCharts({
|
|||
() => buildYDomain(lossData.flatMap((p) => [p.displayLoss, p.displaySmoothed])),
|
||||
[lossData],
|
||||
);
|
||||
const lrDomain = useMemo(
|
||||
() => buildYDomain(lrData.map((p) => p.displayLr)),
|
||||
[lrData],
|
||||
const gradDomain = useMemo(
|
||||
() => buildYDomain(gradData.map((p) => p.displayGradNorm)),
|
||||
[gradData],
|
||||
);
|
||||
const gradNormDomain = useMemo(
|
||||
() => buildYDomain(gradNormData.map((p) => p.displayGradNorm)),
|
||||
[gradNormData],
|
||||
);
|
||||
|
||||
const avgRaw =
|
||||
lossItems.length > 0
|
||||
? +(lossItems.reduce((s, p) => s + p.loss, 0) / lossItems.length).toFixed(4)
|
||||
: 0;
|
||||
|
||||
if (lossItems.length === 0 && lrData.length === 0) return null;
|
||||
|
||||
return (
|
||||
<div className="grid grid-cols-1 gap-4 lg:grid-cols-2">
|
||||
<div className="flex flex-col gap-1">
|
||||
|
|
@ -152,22 +131,13 @@ export function DiffusionCharts({
|
|||
the smoothed line for the trend, not the raw jitter.
|
||||
</p>
|
||||
</div>
|
||||
<LearningRateChartCard
|
||||
data={lrData}
|
||||
domain={lrDomain}
|
||||
<GradNormChartCard
|
||||
data={gradData}
|
||||
domain={gradDomain}
|
||||
visibleStepDomain={stepDomain}
|
||||
xAxisTicks={xAxisTicks}
|
||||
scale="linear"
|
||||
/>
|
||||
{gradNormData.length > 0 && (
|
||||
<GradNormChartCard
|
||||
data={gradNormData}
|
||||
domain={gradNormDomain}
|
||||
visibleStepDomain={stepDomain}
|
||||
xAxisTicks={xAxisTicks}
|
||||
scale="linear"
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
Loading…
Add table
Add a link
Reference in a new issue