Retain diffusion training loss history and expose it in status
The training service kept only the latest loss, so a live loss chart could show a single point. Fold each progress event into bounded (step, loss, lr) history arrays (capped at 4000 points, decimated when full) plus the latest throughput and peak VRAM, and record the family / base model / catalog path on completion. The status endpoint returns these as a nested metric_history object the UI can chart directly, and the start request accepts an optional model_family override.
This commit is contained in:
Daniel Han
parent
7f0a9ebd2f
commit
76520bb553
3 changed files with 101 additions and 2 deletions
|
|
@ -50,6 +50,12 @@ def _default_target(*, event_queue: Any, stop_queue: Any, config: dict) -> None:
|
|||
)
|
||||
|
||||
|
||||
# Cap on retained metric points. When exceeded, the arrays are decimated (every other
|
||||
# point dropped) so a long run stays bounded in memory while the live loss chart keeps a
|
||||
# faithful shape. 4000 points comfortably covers a typical run at full resolution.
|
||||
_METRIC_CAP = 4000
|
||||
|
||||
|
||||
def _idle_state() -> dict[str, Any]:
|
||||
return {
|
||||
"active": False,
|
||||
|
|
@ -65,11 +71,57 @@ def _idle_state() -> dict[str, Any]:
|
|||
"in_model_load": False,
|
||||
"output_dir": None,
|
||||
"lora_path": None,
|
||||
"catalog_path": None,
|
||||
"family": None,
|
||||
"base_model": None,
|
||||
"samples_per_second": None,
|
||||
"peak_memory_gb": None,
|
||||
"started_at": None,
|
||||
"updated_at": None,
|
||||
# Bounded, paired history arrays for the live loss chart (see _append_metric).
|
||||
"metric_steps": [],
|
||||
"metric_loss": [],
|
||||
"metric_lr": [],
|
||||
}
|
||||
|
||||
|
||||
def _append_metric(state: dict[str, Any], step: Any, loss: Any, lr: Any) -> None:
|
||||
"""Append one (step, loss, lr) point to the bounded history arrays on ``state``.
|
||||
|
||||
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 may be None (kept as None so the LR series can be sparse)."""
|
||||
try:
|
||||
istep = int(step)
|
||||
except (TypeError, ValueError):
|
||||
return
|
||||
if istep <= 0 or loss is None:
|
||||
return
|
||||
try:
|
||||
floss = float(loss)
|
||||
except (TypeError, ValueError):
|
||||
return
|
||||
if floss != floss: # NaN guard
|
||||
return
|
||||
flr: Optional[float]
|
||||
try:
|
||||
flr = float(lr) if lr is not None else None
|
||||
except (TypeError, ValueError):
|
||||
flr = None
|
||||
steps = state["metric_steps"]
|
||||
losses = state["metric_loss"]
|
||||
lrs = state["metric_lr"]
|
||||
if len(steps) >= _METRIC_CAP:
|
||||
state["metric_steps"] = steps[::2]
|
||||
state["metric_loss"] = losses[::2]
|
||||
state["metric_lr"] = lrs[::2]
|
||||
steps, losses, lrs = state["metric_steps"], state["metric_loss"], state["metric_lr"]
|
||||
steps.append(istep)
|
||||
losses.append(floss)
|
||||
lrs.append(flr)
|
||||
|
||||
|
||||
class DiffusionTrainingService:
|
||||
"""One diffusion LoRA training job at a time, spawned as a subprocess."""
|
||||
|
||||
|
|
@ -144,6 +196,7 @@ class DiffusionTrainingService:
|
|||
job_id = job_id,
|
||||
status = "running",
|
||||
message = "Starting diffusion LoRA training...",
|
||||
base_model = config.get("base_model") or config.get("model_name"),
|
||||
started_at = now,
|
||||
updated_at = now,
|
||||
)
|
||||
|
|
@ -237,6 +290,14 @@ class DiffusionTrainingService:
|
|||
learning_rate = ev.get("learning_rate", s["learning_rate"]),
|
||||
message = "Training...",
|
||||
)
|
||||
# Fold optional perf fields (emitted by the trainers) so the UI can show
|
||||
# throughput + peak VRAM without a separate channel.
|
||||
if ev.get("samples_per_second") is not None:
|
||||
s["samples_per_second"] = ev.get("samples_per_second")
|
||||
if ev.get("peak_memory_gb") is not None:
|
||||
s["peak_memory_gb"] = ev.get("peak_memory_gb")
|
||||
# Retain a bounded (step, loss, lr) history for the live loss chart.
|
||||
_append_metric(s, ev.get("step"), ev.get("loss"), ev.get("learning_rate"))
|
||||
elif etype == "complete":
|
||||
# Reset in_model_load: a stop during model load emits complete without a
|
||||
# preceding model_load_completed, which would otherwise leave a stale
|
||||
|
|
@ -251,6 +312,12 @@ class DiffusionTrainingService:
|
|||
if ev.get("stopped")
|
||||
else "Training complete.",
|
||||
)
|
||||
if ev.get("catalog_path") is not None:
|
||||
s["catalog_path"] = ev.get("catalog_path")
|
||||
if ev.get("family") is not None:
|
||||
s["family"] = ev.get("family")
|
||||
if ev.get("base_model") is not None:
|
||||
s["base_model"] = ev.get("base_model")
|
||||
elif etype == "error":
|
||||
# Reset in_model_load too: an error raised during model loading has no
|
||||
# model_load_completed, so the terminal state must clear it explicitly.
|
||||
|
|
|
|||
|
|
@ -677,9 +677,13 @@ class DiffusionTrainingStartRequest(BaseModel):
|
|||
|
||||
model_config = ConfigDict(protected_namespaces = ())
|
||||
|
||||
base_model: str = Field(..., description = "HF repo id or local path to an SDXL pipeline")
|
||||
base_model: str = Field(..., description = "HF repo id or local path to a trainable base")
|
||||
data_dir: str = Field(..., description = "Folder of training images (+ captions)")
|
||||
output_dir: str = Field(..., description = "Directory to write the LoRA .safetensors into")
|
||||
model_family: Optional[str] = Field(
|
||||
None,
|
||||
description = "Explicit trainer family (sdxl / flux.1 / ...); omitted = detect from base_model",
|
||||
)
|
||||
instance_prompt: Optional[str] = Field(
|
||||
None, description = "Dreambooth caption applied to images without their own caption"
|
||||
)
|
||||
|
|
@ -720,6 +724,15 @@ class DiffusionTrainingStartResponse(BaseModel):
|
|||
status: str
|
||||
|
||||
|
||||
class DiffusionMetricHistory(BaseModel):
|
||||
"""Paired step-indexed history arrays for the live training charts. ``lr`` entries may
|
||||
be null so a sparse learning-rate series still aligns with ``steps`` by index."""
|
||||
|
||||
steps: List[int] = Field(default_factory = list)
|
||||
loss: List[float] = Field(default_factory = list)
|
||||
lr: List[Optional[float]] = Field(default_factory = list)
|
||||
|
||||
|
||||
class DiffusionTrainingStatusResponse(BaseModel):
|
||||
"""A snapshot of the current diffusion training job (or idle)."""
|
||||
|
||||
|
|
@ -736,8 +749,18 @@ class DiffusionTrainingStatusResponse(BaseModel):
|
|||
in_model_load: bool = False
|
||||
output_dir: Optional[str] = None
|
||||
lora_path: Optional[str] = None
|
||||
# Where the trained adapter was mirrored into the Studio LoRA catalog, and what family
|
||||
# / base it was trained from -- lets the UI deploy the adapter onto the right base.
|
||||
catalog_path: Optional[str] = None
|
||||
family: Optional[str] = None
|
||||
base_model: Optional[str] = None
|
||||
# Live throughput + peak VRAM (from the trainer's progress events).
|
||||
samples_per_second: Optional[float] = None
|
||||
peak_memory_gb: Optional[float] = None
|
||||
started_at: Optional[float] = None
|
||||
updated_at: Optional[float] = None
|
||||
# Bounded step/loss/lr history for the live loss + LR charts.
|
||||
metric_history: Optional[DiffusionMetricHistory] = None
|
||||
|
||||
|
||||
class DiffusionDatasetSummary(BaseModel):
|
||||
|
|
|
|||
|
|
@ -60,6 +60,7 @@ from models import (
|
|||
from models.training import (
|
||||
DiffusionDatasetSummary,
|
||||
DiffusionDatasetUploadResponse,
|
||||
DiffusionMetricHistory,
|
||||
DiffusionTrainingInfoResponse,
|
||||
DiffusionTrainingStartRequest,
|
||||
DiffusionTrainingStartResponse,
|
||||
|
|
@ -1204,7 +1205,15 @@ async def stop_diffusion_training(current_subject: str = Depends(get_current_sub
|
|||
async def diffusion_training_status(current_subject: str = Depends(get_current_subject)):
|
||||
"""Poll the current diffusion training job's status/progress (JSON)."""
|
||||
from core.training.diffusion_training_service import get_diffusion_training_service
|
||||
return DiffusionTrainingStatusResponse(**get_diffusion_training_service().status())
|
||||
|
||||
snap = get_diffusion_training_service().status()
|
||||
# Fold the service's flat history arrays into the nested metric_history the UI charts.
|
||||
metric_history = DiffusionMetricHistory(
|
||||
steps = snap.pop("metric_steps", []),
|
||||
loss = snap.pop("metric_loss", []),
|
||||
lr = snap.pop("metric_lr", []),
|
||||
)
|
||||
return DiffusionTrainingStatusResponse(**snap, metric_history = metric_history)
|
||||
|
||||
|
||||
# Extensions accepted into an image-training dataset folder: images the trainer reads,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue