From 76520bb55393994f716288908a239aa66770a50e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 14:55:17 +0000 Subject: [PATCH] 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. --- .../training/diffusion_training_service.py | 67 +++++++++++++++++++ studio/backend/models/training.py | 25 ++++++- studio/backend/routes/training.py | 11 ++- 3 files changed, 101 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index d0b3d3220e..2ad30e8347 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -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. diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index ec8900c4a3..8385bdebb2 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -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): diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index c56a5e9bc9..c8741785fa 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -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,