From 906f541ac91f075a456cb0146c0b7b6a080f57a3 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 09:49:23 +0000 Subject: [PATCH 01/24] Restructure the diffusion Train tab: right-docked Advanced panel, grayed pre-run charts, confirm-stop with save choice --- studio/frontend/src/features/images/api.ts | 32 +- .../images/train/diffusion-charts.tsx | 8 +- .../images/train/diffusion-train-panel.tsx | 574 ++++++++++++++---- 3 files changed, 476 insertions(+), 138 deletions(-) diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index 911a2ff426..bd6c09dd33 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -281,6 +281,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; } @@ -334,8 +347,16 @@ 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 }), + }), + ); } export async function getDiffusionTrainingStatus(): Promise { @@ -368,6 +389,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. diff --git a/studio/frontend/src/features/images/train/diffusion-charts.tsx b/studio/frontend/src/features/images/train/diffusion-charts.tsx index c4378156cf..daf762f056 100644 --- a/studio/frontend/src/features/images/train/diffusion-charts.tsx +++ b/studio/frontend/src/features/images/train/diffusion-charts.tsx @@ -44,14 +44,16 @@ function fullStepDomain(steps: number[]): [number, number] { } // A diffusion-only metrics view: just Training Loss and Learning Rate, side by side, with a -// note under the loss card explaining why per-step loss looks noisy. +// note under the loss card explaining why per-step loss looks noisy. Always renders both +// cards (even with no data) so the Train tab can show them grayed before a run starts; the +// parent applies the grayed treatment via a wrapper, so we never early-return null here. export function DiffusionCharts({ lossHistory, lrHistory, }: { lossHistory: TrainingSeriesPoint[]; lrHistory: TrainingSeriesPoint[]; -}): ReactElement | null { +}): ReactElement { const lossItems = useMemo(() => toLossItems(lossHistory), [lossHistory]); const smoothed = useMemo( () => (lossItems.length > 0 ? ema(lossItems, SMOOTHING) : []), @@ -109,8 +111,6 @@ export function DiffusionCharts({ ? +(lossItems.reduce((s, p) => s + p.loss, 0) / lossItems.length).toFixed(4) : 0; - if (lossItems.length === 0 && lrData.length === 0) return null; - return (
diff --git a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx index 45df1d50a6..a0fc88fca4 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -3,9 +3,19 @@ import { useCallback, useEffect, useMemo, useRef, useState } from "react"; -import { ArrowDown01Icon } from "@hugeicons/core-free-icons"; +import { LayoutAlignRightIcon, Settings02Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; +import { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, +} from "@/components/ui/alert-dialog"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; @@ -156,6 +166,29 @@ export function DiffusionTrainPanel({ () => families.find((f) => f.name === familyName) ?? families[0], [families, familyName], ); + // The raw backend family record (precision_modes / recommended_precision / supports_compile + // live only here, not on the preset). Absent on an older backend -> the DiT speed controls + // fall back to a sensible default list. + const reportedFamily = useMemo( + () => info?.families?.find((f) => f.name === familyName), + [info?.families, familyName], + ); + // sdxl trains the U-Net in mixed precision (no quantised base), so it uses the + // mixed_precision control instead of base_precision. Everything else is a DiT family. + const isDiT = familyName !== "sdxl"; + // The quantised base precisions this family can train in, with a stable fallback when the + // backend does not report them (older backend, or a preset-only family). + const precisionModes = useMemo>(() => { + const reported = reportedFamily?.precision_modes?.filter( + (m): m is "nf4" | "bf16" | "int8" | "fp8" => + m === "nf4" || m === "bf16" || m === "int8" || m === "fp8", + ); + if (reported && reported.length > 0) return ["auto", ...reported]; + return ["auto", "nf4", "bf16", "int8", "fp8"]; + }, [reportedFamily?.precision_modes]); + // Whether to show the torch.compile control. Default on for DiT families when the backend + // does not say otherwise; sdxl's U-Net path does not expose it here. + const supportsCompile = isDiT && (reportedFamily?.supports_compile ?? true); const [baseChoice, setBaseChoice] = useState(family?.base_repos[0] ?? ""); const [customBase, setCustomBase] = useState(""); @@ -172,19 +205,52 @@ export function DiffusionTrainPanel({ const [outputDir, setOutputDir] = useState(""); const [instancePrompt, setInstancePrompt] = useState(""); - const [showAdvanced, setShowAdvanced] = useState(false); + // The right-docked Advanced settings panel (mirrors the Create tab's / Chat's settings + // panel). Open by default -- on a training tab the hyperparameters are primary content, + // like the LLM Train page's right rail; the fixed top-right button hides it. + const [advancedOpen, setAdvancedOpen] = useState(true); const [steps, setSteps] = useState(500); const [learningRate, setLearningRate] = useState(family?.defaults.lr ?? 0.0001); const [rank, setRank] = useState(family?.defaults.rank ?? 16); const [resolution, setResolution] = useState(family?.defaults.resolution ?? 768); const [batchSize, setBatchSize] = useState(1); + const [gradAccum, setGradAccum] = useState(1); + const [seed, setSeed] = useState(42); + // LR schedule (PR E wired get_scheduler into the loop, so the LR chart reflects this). + // Warmup only applies to the non-constant schedules; plain "constant" ignores it. + const [lrScheduler, setLrScheduler] = useState< + "constant" | "constant_with_warmup" | "cosine" | "linear" + >("constant"); + const [lrWarmupSteps, setLrWarmupSteps] = useState(0); + // Gradient checkpointing trades ~20-30% step time for a large activation-VRAM saving. + const [gradCheckpoint, setGradCheckpoint] = useState(true); + // sdxl (U-Net) trains in a mixed-precision autocast; the DiT families quantise the frozen + // base weights instead (base_precision) and ignore this. Both are surfaced in Advanced. const [precision, setPrecision] = useState<"bf16" | "fp16" | "no">("bf16"); + // Quantised base precision for DiT families (nf4 QLoRA default, or a speed tier). "auto" + // lets the backend pick the family's recommended mode. Re-seeded to the family's + // recommendation on family change (unless the user picked one). + const [basePrecision, setBasePrecision] = useState< + "nf4" | "bf16" | "int8" | "fp8" | "auto" + >("auto"); + // Whether to torch.compile the DiT transformer. "auto" defers to the backend. + const [compileTransformer, setCompileTransformer] = useState<"off" | "on" | "auto">( + "auto", + ); // Track whether the user hand-edited the numeric settings; if not, a family change // re-seeds them from that family's defaults. const settingsDirty = useRef(false); + // Track whether the user hand-picked a base precision; if not, a family change re-seeds it + // from that family's recommended_precision. + const precisionDirty = useRef(false); const [starting, setStarting] = useState(false); const [status, setStatus] = useState(null); + // The confirm-stop dialog (mirrors the LLM Train tab): Continue / Stop / Stop and save. + const [stopDialogOpen, setStopDialogOpen] = useState(false); + // Set when the user confirms a stop; the button reads "Stopping..." until the run ends. + // Clamped to the running state at read time (below) so a fresh run never inherits it. + const [stopRequestedLocal, setStopRequestedLocal] = useState(false); const refreshInfo = useCallback(async (): Promise => { try { @@ -286,7 +352,17 @@ export function DiffusionTrainPanel({ setRank(family.defaults.rank); setResolution(family.defaults.resolution); } - }, [family, loadedBaseRepo]); + // Re-seed the DiT base precision from the family's recommendation (unless the user picked + // one). "auto" is always a safe default when the backend has no recommendation. + if (!precisionDirty.current) { + const rec = reportedFamily?.recommended_precision; + setBasePrecision( + rec === "nf4" || rec === "bf16" || rec === "int8" || rec === "fp8" + ? rec + : "auto", + ); + } + }, [family, loadedBaseRepo, reportedFamily?.recommended_precision]); // The base actually used everywhere (request, deploy, select value). baseChoice can // briefly hold another family's repo between a family switch and the reseed effect @@ -325,6 +401,18 @@ export function DiffusionTrainPanel({ ? Math.min(100, Math.round((status.step / status.total_steps) * 100)) : 0; + // The pending-stop flag only matters while a run is active; clamping at read time (rather + // than resetting in an effect) means a fresh run never inherits a stale "Stopping..." state. + const stopRequested = running && stopRequestedLocal; + + // Whether there is a run to show live (running or a not-yet-dismissed completed run). + // When false, the progress card + charts still render but grayed, as a preview. + const hasRun = Boolean( + status && + status.status !== "idle" && + !(status.status === "completed" && status.job_id === dismissedJobId), + ); + // Notify the parent exactly once per completed run so it rescans the LoRA picker. const notifiedComplete = useRef(false); useEffect(() => { @@ -416,7 +504,9 @@ export function DiffusionTrainPanel({ return toast.error("Resolution must be a multiple of 8 and at least 64."); } if (batchSize < 1) return toast.error("Batch size must be at least 1."); + if (gradAccum < 1) return toast.error("Gradient accumulation must be at least 1."); if (learningRate <= 0) return toast.error("Learning rate must be greater than 0."); + if (lrWarmupSteps < 0) return toast.error("Warmup steps cannot be negative."); setStarting(true); try { await startDiffusionTraining({ @@ -429,8 +519,17 @@ export function DiffusionTrainPanel({ train_steps: steps, learning_rate: learningRate, train_batch_size: batchSize, + gradient_accumulation_steps: gradAccum, + seed, + gradient_checkpointing: gradCheckpoint, + lr_scheduler: lrScheduler, + lr_warmup_steps: lrScheduler === "constant" ? 0 : lrWarmupSteps, lora_rank: rank, mixed_precision: precision, + // DiT families quantise the base weights (base_precision); sdxl uses mixed_precision + // above and ignores this. Only send compile for families that support it. + base_precision: isDiT ? basePrecision : undefined, + compile_transformer: supportsCompile ? compileTransformer : undefined, hf_token: hfApiToken(getHfToken()) || undefined, }); toast.success("Training started"); @@ -452,20 +551,42 @@ export function DiffusionTrainPanel({ steps, learningRate, batchSize, + gradAccum, + seed, + gradCheckpoint, + lrScheduler, + lrWarmupSteps, rank, precision, + isDiT, + basePrecision, + supportsCompile, + compileTransformer, poll, ]); - const onStop = useCallback(async () => { - try { - await stopDiffusionTraining(); - toast.success("Stop requested; finishing the current step."); - void poll(); - } catch (e) { - toast.error(e instanceof Error ? e.message : "Failed to stop training"); - } - }, [poll]); + // Confirm-then-stop, mirroring the LLM Train tab. `save` writes the current adapter before + // halting ("Stop and save"); false discards it ("Stop"). Closes the dialog and marks the + // stop as requested so the button reads "Stopping..." until the backend reports it stopped. + const onStop = useCallback( + async (save: boolean) => { + setStopDialogOpen(false); + setStopRequestedLocal(true); + try { + await stopDiffusionTraining(save); + toast.success( + save + ? "Stop requested; saving the adapter after the current step." + : "Stop requested; discarding this run after the current step.", + ); + void poll(); + } catch (e) { + setStopRequestedLocal(false); + toast.error(e instanceof Error ? e.message : "Failed to stop training"); + } + }, + [poll], + ); const onDeployClick = useCallback(() => { if (!status?.catalog_path) { @@ -508,6 +629,135 @@ export function DiffusionTrainPanel({
); + const precisionLabel = (m: "nf4" | "bf16" | "int8" | "fp8" | "auto"): string => { + if (m === "auto") return "Auto (recommended)"; + if (m === "nf4") return "nf4 (4-bit QLoRA, lowest VRAM)"; + if (m === "bf16") return "bf16 (fastest, most VRAM)"; + if (m === "int8") return "int8 (8-bit)"; + return "fp8 (experimental)"; + }; + + // The Advanced training settings, rendered inside the right-docked panel (mirrors the + // Create tab's advancedControls). Numeric hyperparameters, the sdxl/DiT precision control, + // and the DiT speed levers (base precision + torch.compile). + const advancedControls = ( +
+
+ {numberField("Steps", steps, setSteps, 1)} + {numberField("LoRA rank", rank, setRank, 1)} + {numberField("Resolution", resolution, setResolution, 512, { min: 64, step: 64 })} + {numberField("Batch", batchSize, setBatchSize, 1)} + {numberField("Grad accumulation", gradAccum, setGradAccum, 1)} + {numberField("Seed", seed, setSeed, 42, { min: 0 })} +
+ {numberField("Learning rate", learningRate, setLearningRate, 0.0001, { + min: 0, + step: 0.00001, + })} + +
+ + + {lrScheduler !== "constant" && + numberField("Warmup steps", lrWarmupSteps, setLrWarmupSteps, 0, { min: 0 })} +

+ How the learning rate evolves over the run (shown live in the LR chart). +

+
+ +
+ + +

+ Recomputes activations in the backward pass: a large VRAM saving for a modest + per-step slowdown. +

+
+ + {isDiT ? ( + <> +
+ + +

+ How the frozen base weights are quantised. nf4 (4-bit) uses the least VRAM; + bf16 is fastest but needs the most. Auto picks this family's recommended mode. +

+
+ {supportsCompile && ( +
+ + +

+ torch.compile the transformer. Adds a one-time warmup, then speeds up each step. +

+
+ )} + + ) : ( +
+ + +

+ Mixed-precision autocast for the U-Net. bf16 suits modern GPUs. +

+
+ )} +
+ ); + return (
{/* Left: configure */} @@ -714,54 +964,28 @@ export function DiffusionTrainPanel({ />
- {/* Collapsed training settings */} - - {showAdvanced && ( - <> -
- {numberField("Steps", steps, setSteps, 1)} - {numberField("LoRA rank", rank, setRank, 1)} - {numberField("Resolution", resolution, setResolution, 512, { min: 64, step: 64 })} - {numberField("Batch", batchSize, setBatchSize, 1)} -
-
- {numberField("Learning rate", learningRate, setLearningRate, 0.0001, { - min: 0, - step: 0.00001, - })} -
- - -
-
- - )} + {advancedOpen + ? "Training settings are in the Advanced panel." + : "Adjust steps, rank, precision and speed in Advanced settings."} +
{running ? ( - ) : (
- {/* Right: run view */} -
- {status && - status.status !== "idle" && - !(status.status === "completed" && status.job_id === dismissedJobId) ? ( - <> -
-
- {status.status} - - {status.total_steps > 0 ? `${status.step}/${status.total_steps} steps` : ""} - -
-
-
-
-
- - - - -
- {status.message && ( -

{status.message}

- )} -
- - - - {completed && ( -
- Adapter ready -

- Trained{status.family ? ` (${status.family})` : ""} and added to the LoRA - picker. - {status.lora_path && ( - Saved: {status.lora_path} - )} -

-
- - -
-
+ {/* Right: run view. The progress card + charts are ALWAYS mounted; before a run they + render grayed as a preview (with an overlaid hint), so the layout never jumps when + training starts. A single fixed top-right button toggles the Advanced panel. */} +
+ {/* Fixed Advanced toggle (mirrors Chat's / the Create tab's settings toggle: same icon + in both states so it never moves, highlighted when open). */} +
+ +
+ +
+
+
+ + {hasRun ? status?.status : "Idle"} + + + {hasRun && (status?.total_steps ?? 0) > 0 + ? `${status?.step}/${status?.total_steps} steps` + : ""} + +
+
+
+
+
+ + + + +
+ {hasRun && status?.message && ( +

{status.message}

+ )} +
+ + +
+ + {/* Placeholder hint overlaid on the grayed preview until a run exists. */} + {!hasRun && ( +
+

No training run yet

- Pick a family and dataset on the left, then Start training. The loss chart and - progress appear here live. + Pick a family and dataset on the left, then Start training. Progress and the loss + chart fill in here live.

)} + + {completed && ( +
+ Adapter ready +

+ Trained{status?.family ? ` (${status.family})` : ""} and added to the LoRA + picker. + {status?.lora_path && ( + Saved: {status.lora_path} + )} +

+
+ + +
+
+ )}
+ + {/* Right-docked Advanced settings panel (mirrors the Create tab / Chat settings panel): + closed by default, opened by the top-right toggle above. Holds every training + hyperparameter and the DiT speed levers so the left rail stays focused on data. */} + {advancedOpen && ( +
+
+ + + Advanced + + +
+
+

+ Defaults suit a first run. Changes apply to the next Start training. +

+ {advancedControls} +
+
+ )} + + {/* Confirm-stop dialog (mirrors the LLM Train tab): Continue / Stop / Stop and save. */} + + + + Stop training? + + Save the adapter trained so far, or discard this run? Either way the current step + finishes first. + + + + Continue training + void onStop(false)}> + Stop + + void onStop(true)}> + Stop and save + + + +
); } From 7da9396a126af0f3d112a0d31180fe52e47993b9 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 09:54:08 +0000 Subject: [PATCH 02/24] Surface the partial adapter after Stop and save: deploy card for stopped runs with a saved LoRA --- .../images/train/diffusion-train-panel.tsx | 28 ++++++++++++++----- 1 file changed, 21 insertions(+), 7 deletions(-) diff --git a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx index a0fc88fca4..5f37c53b27 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -396,6 +396,12 @@ export function DiffusionTrainPanel({ const running = Boolean(status?.active) || status?.status === "running"; const completed = status?.status === "completed" && status.job_id !== dismissedJobId; + // "Stop and save" ends the run as "stopped" WITH a saved partial adapter; it must get + // the same ready-to-deploy card as a full run (only a no-save stop has nothing to show). + const stoppedWithAdapter = + status?.status === "stopped" && + Boolean(status?.lora_path) && + status.job_id !== dismissedJobId; const pct = status && status.total_steps > 0 ? Math.min(100, Math.round((status.step / status.total_steps) * 100)) @@ -413,16 +419,20 @@ export function DiffusionTrainPanel({ !(status.status === "completed" && status.job_id === dismissedJobId), ); - // Notify the parent exactly once per completed run so it rescans the LoRA picker. + // Notify the parent exactly once per run that produced an adapter (full completion or + // stop-and-save) so it rescans the LoRA picker. const notifiedComplete = useRef(false); useEffect(() => { - if (status?.status === "completed" && !notifiedComplete.current) { + const producedAdapter = + status?.status === "completed" || + (status?.status === "stopped" && Boolean(status?.lora_path)); + if (producedAdapter && !notifiedComplete.current) { notifiedComplete.current = true; onTrainingComplete?.(); } else if (status?.status === "running" && notifiedComplete.current) { notifiedComplete.current = false; } - }, [status?.status, onTrainingComplete]); + }, [status?.status, status?.lora_path, onTrainingComplete]); const selectedDataset = dataset !== UPLOAD_DATASET ? info?.datasets.find((d) => d.name === dataset) : undefined; @@ -1095,12 +1105,16 @@ export function DiffusionTrainPanel({
)} - {completed && ( + {(completed || stoppedWithAdapter) && (
- Adapter ready + + {completed ? "Adapter ready" : "Partial adapter saved"} +

- Trained{status?.family ? ` (${status.family})` : ""} and added to the LoRA - picker. + {completed + ? "Trained" + : "Stopped early; the adapter as of the last finished step was saved"} + {status?.family ? ` (${status.family})` : ""} and added to the LoRA picker. {status?.lora_path && ( Saved: {status.lora_path} )} From 923c1c0d8085f2a42deb59d751dd617cd982bd6c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 10:03:04 +0000 Subject: [PATCH 03/24] Reset the pending-stop flag when a new run starts (stale flag disabled the Stop button on the next run) --- .../src/features/images/train/diffusion-train-panel.tsx | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx index 5f37c53b27..9eebcf9bb6 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -518,6 +518,10 @@ export function DiffusionTrainPanel({ if (learningRate <= 0) return toast.error("Learning rate must be greater than 0."); if (lrWarmupSteps < 0) return toast.error("Warmup steps cannot be negative."); setStarting(true); + // A previous run's confirmed stop must not leak into this run: without the reset the + // read-time clamp (running && stopRequestedLocal) re-arms the moment the new run goes + // active, rendering a permanently disabled "Stopping..." button. + setStopRequestedLocal(false); try { await startDiffusionTraining({ base_model: baseModel, From b9f2a71f63b02738d1d4d8e2b30e7e1fb6c039de Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 11:05:58 +0000 Subject: [PATCH 04/24] Train tab reflow: tabs on the left, settings as the run area until training starts --- .../src/features/images/images-page.tsx | 31 +- .../images/train/diffusion-train-panel.tsx | 404 ++++++++---------- 2 files changed, 189 insertions(+), 246 deletions(-) diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 8d8f0a38f3..bff5f93e9d 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -1956,21 +1956,22 @@ export function ImagesPage({ active = true }: { active?: boolean }) { shared element matches. The load progress shows in a chat-style toast, not here. ── */}

- setSelectorOpen(active && o)} - />
- {/* Create | Train page-mode switch, next to the model selector. Create is the - generation workspace; Train is the full-page LoRA training workspace. */} + 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. */} setPageMode(v as "create" | "train")}> @@ -1981,6 +1982,8 @@ export function ImagesPage({ active = true }: { active?: boolean }) { +
+
{/* 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. */} diff --git a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx index 9eebcf9bb6..5a38d24910 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -3,7 +3,7 @@ import { useCallback, useEffect, useMemo, useRef, useState } from "react"; -import { LayoutAlignRightIcon, Settings02Icon } from "@hugeicons/core-free-icons"; +import { Settings02Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { @@ -205,10 +205,6 @@ export function DiffusionTrainPanel({ const [outputDir, setOutputDir] = useState(""); const [instancePrompt, setInstancePrompt] = useState(""); - // The right-docked Advanced settings panel (mirrors the Create tab's / Chat's settings - // panel). Open by default -- on a training tab the hyperparameters are primary content, - // like the LLM Train page's right rail; the fixed top-right button hides it. - const [advancedOpen, setAdvancedOpen] = useState(true); const [steps, setSteps] = useState(500); const [learningRate, setLearningRate] = useState(family?.defaults.lr ?? 0.0001); const [rank, setRank] = useState(family?.defaults.rank ?? 16); @@ -651,12 +647,12 @@ export function DiffusionTrainPanel({ return "fp8 (experimental)"; }; - // The Advanced training settings, rendered inside the right-docked panel (mirrors the - // Create tab's advancedControls). Numeric hyperparameters, the sdxl/DiT precision control, - // and the DiT speed levers (base precision + torch.compile). - const advancedControls = ( -
-
+ // The training settings, shown as the run area's MAIN content before a run starts + // (settings are set once, up front); once training starts the run view (progress + + // charts) replaces them. Laid out as a wide grid for the center column. + const trainingSettings = ( +
+
{numberField("Steps", steps, setSteps, 1)} {numberField("LoRA rank", rank, setRank, 1)} {numberField("Resolution", resolution, setResolution, 512, { min: 64, step: 64 })} @@ -664,111 +660,117 @@ export function DiffusionTrainPanel({ {numberField("Grad accumulation", gradAccum, setGradAccum, 1)} {numberField("Seed", seed, setSeed, 42, { min: 0 })}
- {numberField("Learning rate", learningRate, setLearningRate, 0.0001, { - min: 0, - step: 0.00001, - })} -
- - - {lrScheduler !== "constant" && - numberField("Warmup steps", lrWarmupSteps, setLrWarmupSteps, 0, { min: 0 })} -

- How the learning rate evolves over the run (shown live in the LR chart). -

-
- -
- - -

- Recomputes activations in the backward pass: a large VRAM saving for a modest - per-step slowdown. -

-
- - {isDiT ? ( - <> -
- - -

- How the frozen base weights are quantised. nf4 (4-bit) uses the least VRAM; - bf16 is fastest but needs the most. Auto picks this family's recommended mode. -

-
- {supportsCompile && ( -
- - -

- torch.compile the transformer. Adds a one-time warmup, then speeds up each step. -

-
- )} - - ) : ( +
+ {numberField("Learning rate", learningRate, setLearningRate, 0.0001, { + min: 0, + step: 0.00001, + })}
- +

- Mixed-precision autocast for the U-Net. bf16 suits modern GPUs. + How the learning rate evolves over the run (shown live in the LR chart).

- )} + {lrScheduler !== "constant" && + numberField("Warmup steps", lrWarmupSteps, setLrWarmupSteps, 0, { min: 0 })} +
+ +
+
+ + +

+ Recomputes activations in the backward pass: a large VRAM saving for a modest + per-step slowdown. +

+
+ + {isDiT ? ( + <> +
+ + +

+ How the frozen base weights are quantised. nf4 (4-bit) uses the least VRAM; + bf16 is fastest but needs the most. Auto picks this family's recommended + mode. +

+
+ {supportsCompile && ( +
+ + +

+ torch.compile the transformer. Adds a one-time warmup, then speeds up each + step. +

+
+ )} + + ) : ( +
+ + +

+ Mixed-precision autocast for the U-Net. bf16 suits modern GPUs. +

+
+ )} +
); @@ -978,18 +980,6 @@ export function DiffusionTrainPanel({ />
- {/* Training settings now live in the right-docked Advanced panel (opened by the - top-right toggle), so the left rail stays focused on the dataset + trigger. */} - -
{running ? (
- {/* Right: run view. The progress card + charts are ALWAYS mounted; before a run they - render grayed as a preview (with an overlaid hint), so the layout never jumps when - training starts. A single fixed top-right button toggles the Advanced panel. */} + {/* Right: the run area. Before a run it shows the training settings (they are set + once, up front); the moment a run exists the settings give way to the run view + (progress + live charts), and come back after "Train another". */}
- {/* Fixed Advanced toggle (mirrors Chat's / the Create tab's settings toggle: same icon - in both states so it never moves, highlighted when open). */} -
- -
- -
-
+ {!hasRun ? ( +
- - {hasRun ? status?.status : "Idle"} + + + Training settings - {hasRun && (status?.total_steps ?? 0) > 0 - ? `${status?.step}/${status?.total_steps} steps` - : ""} + Applied when you press Start training
-
-
-
-
- - - - -
- {hasRun && status?.message && ( -

{status.message}

- )} + {trainingSettings} +

+ Once training starts, live progress and the Training Loss / Learning Rate + charts take over this area. +

- - -
- - {/* Placeholder hint overlaid on the grayed preview until a run exists. */} - {!hasRun && ( -
-
-

No training run yet

-

- Pick a family and dataset on the left, then Start training. Progress and the loss - chart fill in here live. -

+ ) : ( + <> +
+
+ {status?.status} + + {(status?.total_steps ?? 0) > 0 + ? `${status?.step}/${status?.total_steps} steps` + : ""} + +
+
+
+
+
+ + + + +
+ {status?.message && ( +

{status.message}

+ )}
-
+ + + )} {(completed || stoppedWithAdapter) && ( @@ -1140,34 +1108,6 @@ export function DiffusionTrainPanel({ )}
- {/* Right-docked Advanced settings panel (mirrors the Create tab / Chat settings panel): - closed by default, opened by the top-right toggle above. Holds every training - hyperparameter and the DiT speed levers so the left rail stays focused on data. */} - {advancedOpen && ( -
-
- - - Advanced - - -
-
-

- Defaults suit a first run. Changes apply to the next Start training. -

- {advancedControls} -
-
- )} - {/* Confirm-stop dialog (mirrors the LLM Train tab): Continue / Stop / Stop and save. */} From 0fbdd743a07c3cd287339fac28cd24f2b98b6435 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 11:10:40 +0000 Subject: [PATCH 05/24] Report grad norm from the trainers and chart it instead of LR; celebrate completion in the run header --- .../core/training/diffusion_dit_trainer.py | 6 ++- .../core/training/diffusion_lora_trainer.py | 7 ++- .../training/diffusion_training_service.py | 47 ++++++++++++++----- studio/backend/models/training.py | 9 +++- studio/backend/routes/training.py | 1 + studio/frontend/src/features/images/api.ts | 3 ++ .../images/train/diffusion-charts.tsx | 47 ++++++++++--------- .../images/train/diffusion-train-panel.tsx | 12 +++-- 8 files changed, 89 insertions(+), 43 deletions(-) diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 5221425e4b..5776f87652 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -1157,8 +1157,11 @@ 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 = None if cfg.max_grad_norm and cfg.max_grad_norm > 0: - torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm) + # 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() @@ -1185,6 +1188,7 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto loss = round(step_loss, 5), avg_loss = round(running_loss / done, 5), learning_rate = lr_sched.get_last_lr()[0], + grad_norm = round(grad_norm, 5) if grad_norm is not None else None, samples_per_second = sps, peak_memory_gb = peak_gb or None, ) diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index 7c7ca79d51..10b6cd89cd 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -509,8 +509,12 @@ 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 = None if cfg.max_grad_norm and cfg.max_grad_norm > 0: - torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm) + # 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() @@ -535,6 +539,7 @@ def run_diffusion_lora_training( loss = round(step_loss, 5), avg_loss = round(running_loss / done, 5), learning_rate = lr_sched.get_last_lr()[0], + grad_norm = round(grad_norm, 5) if grad_norm is not None else None, samples_per_second = samples_per_second, peak_memory_gb = peak_gb or None, ) diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index fa44cceabe..2251038b78 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -67,6 +67,7 @@ def _idle_state() -> dict[str, Any]: "loss": None, "avg_loss": None, "learning_rate": None, + "grad_norm": None, "num_images": None, "in_model_load": False, "output_dir": None, @@ -82,16 +83,21 @@ def _idle_state() -> dict[str, Any]: "metric_steps": [], "metric_loss": [], "metric_lr": [], + "metric_grad_norm": [], } -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``. +def _append_metric( + state: dict[str, Any], step: Any, loss: Any, lr: Any, grad_norm: Any = None +) -> None: + """Append one (step, loss, lr, grad_norm) 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).""" + 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): @@ -104,22 +110,34 @@ def _append_metric(state: dict[str, Any], step: Any, loss: Any, lr: Any) -> None 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 + + def _opt_float(v: Any) -> Optional[float]: + try: + return float(v) if v is not None else None + except (TypeError, ValueError): + return None + + flr = _opt_float(lr) + fgn = _opt_float(grad_norm) steps = state["metric_steps"] losses = state["metric_loss"] lrs = state["metric_lr"] + gns = state["metric_grad_norm"] 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"] + state["metric_grad_norm"] = gns[::2] + steps, losses, lrs, gns = ( + state["metric_steps"], + state["metric_loss"], + state["metric_lr"], + state["metric_grad_norm"], + ) steps.append(istep) losses.append(floss) lrs.append(flr) + gns.append(fgn) class DiffusionTrainingService: @@ -315,6 +333,7 @@ class DiffusionTrainingService: loss = ev.get("loss", s["loss"]), avg_loss = ev.get("avg_loss", s["avg_loss"]), learning_rate = ev.get("learning_rate", s["learning_rate"]), + grad_norm = ev.get("grad_norm", s["grad_norm"]), message = "Training...", ) # Fold optional perf fields (emitted by the trainers) so the UI can show @@ -323,8 +342,14 @@ class DiffusionTrainingService: 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")) + # Retain a bounded (step, loss, lr, grad_norm) history for the live charts. + _append_metric( + s, + ev.get("step"), + ev.get("loss"), + ev.get("learning_rate"), + ev.get("grad_norm"), + ) 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 diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 58a9779609..41452ddc89 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -753,12 +753,14 @@ class DiffusionTrainingStartResponse(BaseModel): 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.""" + """Paired step-indexed history arrays for the live training charts. ``lr`` and + ``grad_norm`` entries may be null so those sparse series still align 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) + grad_norm: List[Optional[float]] = Field(default_factory = list) class DiffusionTrainingStatusResponse(BaseModel): @@ -773,6 +775,9 @@ class DiffusionTrainingStatusResponse(BaseModel): loss: Optional[float] = None avg_loss: Optional[float] = None learning_rate: Optional[float] = None + # 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 output_dir: Optional[str] = None diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 271abebd78..b16d95add0 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1270,6 +1270,7 @@ async def diffusion_training_status(current_subject: str = Depends(get_current_s steps = snap.pop("metric_steps", []), loss = snap.pop("metric_loss", []), lr = snap.pop("metric_lr", []), + grad_norm = snap.pop("metric_grad_norm", []), ) return DiffusionTrainingStatusResponse(**snap, metric_history = metric_history) diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index bd6c09dd33..25e4f103e4 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -304,6 +304,8 @@ export interface DiffusionMetricHistory { steps: number[]; loss: number[]; lr: Array; + // Total pre-clip gradient norm per step (the training health signal the charts show). + grad_norm?: Array; } // A snapshot of the current diffusion training job (GET /api/train/diffusion/status). @@ -317,6 +319,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; diff --git a/studio/frontend/src/features/images/train/diffusion-charts.tsx b/studio/frontend/src/features/images/train/diffusion-charts.tsx index daf762f056..7deb70f122 100644 --- a/studio/frontend/src/features/images/train/diffusion-charts.tsx +++ b/studio/frontend/src/features/images/train/diffusion-charts.tsx @@ -4,12 +4,13 @@ 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 { LearningRateChartCard } from "@/features/studio/sections/charts/learning-rate-chart-card"; +import { GradNormChartCard } from "@/features/studio/sections/charts/grad-norm-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 @@ -43,16 +44,16 @@ function fullStepDomain(steps: number[]): [number, number] { return [min, max]; } -// A diffusion-only metrics view: just Training Loss and Learning Rate, 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 Train tab can show them grayed before a run starts; the -// parent applies the grayed treatment via a wrapper, so we never early-return null here. +// 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, }: { lossHistory: TrainingSeriesPoint[]; - lrHistory: TrainingSeriesPoint[]; + gradNormHistory: TrainingSeriesPoint[]; }): ReactElement { const lossItems = useMemo(() => toLossItems(lossHistory), [lossHistory]); const smoothed = useMemo( @@ -73,23 +74,23 @@ export function DiffusionCharts({ [reducedLoss], ); - const lrData = useMemo( + const gradData = useMemo( () => compressSeries( - lrHistory + gradNormHistory .filter((p) => Number.isFinite(p.value)) - .map((p) => ({ step: p.step, lr: p.value, displayLr: p.value })), + .map((p) => ({ step: p.step, gradNorm: p.value, displayGradNorm: p.value })), MAX_RENDER_POINTS, ), - [lrHistory], + [gradNormHistory], ); const steps = useMemo(() => { const set = new Set(); for (const p of lossData) set.add(p.step); - for (const p of lrData) set.add(p.step); + for (const p of gradData) set.add(p.step); return Array.from(set).sort((a, b) => a - b); - }, [lossData, lrData]); + }, [lossData, gradData]); const stepDomain = useMemo(() => fullStepDomain(steps), [steps]); const xAxisTicks = useMemo( @@ -101,9 +102,9 @@ 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 avgRaw = @@ -131,9 +132,9 @@ export function DiffusionCharts({ the smoothed line for the trend, not the raw jitter.

- ({ step, value: h.loss[i] })).filter((p) => p.value != null); }, [status?.metric_history]); - const lrHistory: TrainingSeriesPoint[] = useMemo(() => { + const gradNormHistory: TrainingSeriesPoint[] = useMemo(() => { const h = status?.metric_history; - if (!h) return []; + if (!h?.grad_norm) return []; return h.steps - .map((step, i) => ({ step, value: h.lr[i] })) + .map((step, i) => ({ step, value: h.grad_norm?.[i] ?? null })) .filter((p): p is TrainingSeriesPoint => p.value != null); }, [status?.metric_history]); @@ -1029,7 +1029,9 @@ export function DiffusionTrainPanel({ <>
- {status?.status} + + {status?.status === "completed" ? "Training complete \u{1F389}" : status?.status} + {(status?.total_steps ?? 0) > 0 ? `${status?.step}/${status?.total_steps} steps` @@ -1073,7 +1075,7 @@ export function DiffusionTrainPanel({ )}
- + )} From bc0d10f7597a7aa7fa56fba4732fe99edfe36107 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 3 Jul 2026 11:14:46 +0000 Subject: [PATCH 06/24] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/training/diffusion_lora_trainer.py | 4 +--- studio/backend/core/training/diffusion_training_service.py | 6 +++++- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index 10b6cd89cd..84bbb6bf74 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -512,9 +512,7 @@ def run_diffusion_lora_training( grad_norm = None if cfg.max_grad_norm and cfg.max_grad_norm > 0: # 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) - ) + grad_norm = float(torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm)) optimizer.step() lr_sched.step() diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index 2251038b78..9fbf9fdf25 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -88,7 +88,11 @@ def _idle_state() -> dict[str, Any]: def _append_metric( - state: dict[str, Any], step: Any, loss: Any, lr: Any, grad_norm: Any = None + state: dict[str, Any], + step: Any, + loss: Any, + lr: Any, + grad_norm: Any = None, ) -> None: """Append one (step, loss, lr, grad_norm) point to the bounded history arrays on ``state``. From 32c5855742c871bcc0f1daf23b24b66312fe5a90 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 11:52:25 +0000 Subject: [PATCH 07/24] Train tab: Stop inside the run card, adapter card above charts, persisted run history with re-plottable logs, clearer stop dialog --- .../training/diffusion_training_service.py | 103 ++++++ studio/backend/models/training.py | 37 +++ studio/backend/routes/training.py | 30 ++ .../backend/tests/test_diffusion_training.py | 100 ++++++ studio/frontend/src/features/images/api.ts | 41 +++ .../images/train/diffusion-train-panel.tsx | 297 +++++++++++++----- 6 files changed, 537 insertions(+), 71 deletions(-) diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index 2251038b78..8fef55641e 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -17,10 +17,13 @@ runs a scripted target on a thread. from __future__ import annotations +import json 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 @@ -56,6 +59,54 @@ 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()) + except Exception: # noqa: BLE001 -- a corrupt record never breaks the listing + 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()) + except Exception: # noqa: BLE001 -- missing/corrupt record + return None + + def _idle_state() -> dict[str, Any]: return { "active": False, @@ -156,6 +207,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: @@ -218,6 +271,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 ) @@ -279,12 +334,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)) + except Exception: # noqa: BLE001 -- persisting history must never break the run + pass + def _apply_event( self, ev: dict[str, Any], diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 41452ddc89..d7703f7141 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -796,6 +796,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.""" diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index b16d95add0..a25983dcc3 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -70,6 +70,9 @@ from models.training import ( DiffusionMetricHistory, DiffusionTrainableFamily, DiffusionTrainingInfoResponse, + DiffusionTrainingRunDetail, + DiffusionTrainingRunsResponse, + DiffusionTrainingRunSummary, DiffusionTrainingStartRequest, DiffusionTrainingStartResponse, DiffusionTrainingStatusResponse, @@ -1275,6 +1278,33 @@ 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 + + return DiffusionTrainingRunsResponse( + runs = [DiffusionTrainingRunSummary(**r) for r in list_diffusion_runs(limit = limit)] + ) + + +@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) + if rec is None: + raise HTTPException(status_code = 404, detail = "No such training run.") + return DiffusionTrainingRunDetail(**rec) + + # 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"} diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index f3b7d2b58a..ad43cf85f0 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -65,6 +65,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"}) @@ -625,3 +637,91 @@ 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 diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index 25e4f103e4..61726a02ff 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -362,6 +362,47 @@ export async function stopDiffusionTraining(save = true): Promise<{ status: stri ); } +// 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 | 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 { + return parseJson(await authFetch(`/api/train/diffusion/runs/${encodeURIComponent(jobId)}`)); +} + export async function getDiffusionTrainingStatus(): Promise { return parseJson(await authFetch("/api/train/diffusion/status")); } diff --git a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx index 28da6296d7..07ddd01c1c 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -29,10 +29,14 @@ import { type DiffusionDatasetExample, type DiffusionTrainableFamily, type DiffusionTrainingInfo, + type DiffusionTrainingRunDetail, + type DiffusionTrainingRunSummary, type DiffusionTrainingStatus, getDiffusionTrainingInfo, + getDiffusionTrainingRun, getDiffusionTrainingStatus, listDiffusionDatasetExamples, + listDiffusionTrainingRuns, startDiffusionTraining, stopDiffusionTraining, uploadDiffusionDataset, @@ -242,6 +246,10 @@ export function DiffusionTrainPanel({ const [starting, setStarting] = useState(false); const [status, setStatus] = useState(null); + // Persisted previous runs (terminal), listed on the idle view; selecting one loads its + // full record (config + metric logs) and re-plots its charts read-only. + const [prevRuns, setPrevRuns] = useState([]); + const [viewRun, setViewRun] = useState(null); // The confirm-stop dialog (mirrors the LLM Train tab): Continue / Stop / Stop and save. const [stopDialogOpen, setStopDialogOpen] = useState(false); // Set when the user confirms a stop; the button reads "Stopping..." until the run ends. @@ -454,6 +462,44 @@ export function DiffusionTrainPanel({ .filter((p): p is TrainingSeriesPoint => p.value != null); }, [status?.metric_history]); + // Refresh the previous-runs list whenever the service is not mid-run (on mount and + // right after a run terminates, when its record has just been persisted). + useEffect(() => { + if (!active) return; + if (status?.status === "running") return; + let cancelled = false; + listDiffusionTrainingRuns() + .then((r) => { + if (!cancelled) setPrevRuns(r.runs); + }) + .catch(() => {}); + return () => { + cancelled = true; + }; + }, [active, status?.status]); + + const openPrevRun = useCallback(async (jobId: string) => { + try { + setViewRun(await getDiffusionTrainingRun(jobId)); + } catch (e) { + toast.error(e instanceof Error ? e.message : "Could not load that run"); + } + }, []); + + // Chart series for a selected previous run (from its persisted metric logs). + const viewLossHistory: TrainingSeriesPoint[] = useMemo(() => { + const h = viewRun?.metric_history; + if (!h) return []; + return h.steps.map((step, i) => ({ step, value: h.loss[i] })).filter((p) => p.value != null); + }, [viewRun?.metric_history]); + const viewGradNormHistory: TrainingSeriesPoint[] = useMemo(() => { + const h = viewRun?.metric_history; + if (!h?.grad_norm) return []; + return h.steps + .map((step, i) => ({ step, value: h.grad_norm?.[i] ?? null })) + .filter((p): p is TrainingSeriesPoint => p.value != null); + }, [viewRun?.metric_history]); + const onUpload = useCallback(async () => { const files = Array.from(fileInputRef.current?.files ?? []); if (files.length === 0) { @@ -518,6 +564,8 @@ export function DiffusionTrainPanel({ // read-time clamp (running && stopRequestedLocal) re-arms the moment the new run goes // active, rendering a permanently disabled "Stopping..." button. setStopRequestedLocal(false); + // A history view must not shadow the new live run. + setViewRun(null); try { await startDiffusionTraining({ base_model: baseModel, @@ -980,51 +1028,145 @@ export function DiffusionTrainPanel({ />
+ {/* Start lives here; Stop lives in the run card next to the live stats. */}
- {running ? ( - - ) : ( - - )} +
- {/* Right: the run area. Before a run it shows the training settings (they are set - once, up front); the moment a run exists the settings give way to the run view - (progress + live charts), and come back after "Train another". */} + {/* Right: the run area. Before a run it shows the training settings (set once, up + front) plus the previous-runs history; during/after a run the live view takes + over (progress with Stop, then the saved-adapter card ABOVE the charts). + Selecting a previous run re-plots its persisted logs read-only. */}
- {!hasRun ? ( -
-
- - - Training settings - - - Applied when you press Start training - + {viewRun && !hasRun ? ( + <> +
+
+ + Previous run: {viewRun.adapter || viewRun.job_id.slice(0, 8)} + + +
+
+ + + + +
+

+ {viewRun.family ? `${viewRun.family} - ` : ""} + {viewRun.base_model || ""} + {viewRun.ended_at + ? ` - ${new Date(viewRun.ended_at * 1000).toLocaleString()}` + : ""} +

+ {viewRun.saved && viewRun.catalog_path && ( +
+ +
+ )}
- {trainingSettings} -

- Once training starts, live progress and the Training Loss / Learning Rate - charts take over this area. -

-
+ + + ) : !hasRun ? ( + <> +
+
+ + + Training settings + + + Applied when you press Start training + +
+ {trainingSettings} +

+ Once training starts, live progress and the Training Loss / Gradient Norm + charts take over this area. +

+
+ + {prevRuns.length > 0 && ( +
+ Previous runs +
+ {prevRuns.map((r) => ( + + ))} +
+
+ )} + ) : ( <>
@@ -1073,41 +1215,52 @@ export function DiffusionTrainPanel({ {status?.message && (

{status.message}

)} + {running && ( + + )}
+ {(completed || stoppedWithAdapter) && ( +
+ + {completed ? "Adapter ready" : "Partial adapter saved"} + +

+ {completed + ? "Trained" + : "Stopped early; the adapter as of the last finished step was saved"} + {status?.family ? ` (${status.family})` : ""} and added to the LoRA picker. + {status?.lora_path && ( + Saved: {status.lora_path} + )} +

+
+ + +
+
+ )} + )} - - {(completed || stoppedWithAdapter) && ( -
- - {completed ? "Adapter ready" : "Partial adapter saved"} - -

- {completed - ? "Trained" - : "Stopped early; the adapter as of the last finished step was saved"} - {status?.family ? ` (${status.family})` : ""} and added to the LoRA picker. - {status?.lora_path && ( - Saved: {status.lora_path} - )} -

-
- - -
-
- )}
{/* Confirm-stop dialog (mirrors the LLM Train tab): Continue / Stop / Stop and save. */} @@ -1120,10 +1273,12 @@ export function DiffusionTrainPanel({ finishes first. - + {/* items-center + a real label on the destructive action: a bare "Stop" rendered + as a stubby pill between two wide ones and read as misaligned. */} + Continue training void onStop(false)}> - Stop + Stop without saving void onStop(true)}> Stop and save From 2fd21df1dcc57b5f2f88a2842a1c4c50d72dcbbc Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 11:57:44 +0000 Subject: [PATCH 08/24] Dismiss any terminal run back to settings (stopped and error runs were trapped in the run view) --- .../images/train/diffusion-train-panel.tsx | 25 ++++++++++++++++--- 1 file changed, 22 insertions(+), 3 deletions(-) diff --git a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx index 07ddd01c1c..0a9f44b839 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -415,12 +415,15 @@ export function DiffusionTrainPanel({ // than resetting in an effect) means a fresh run never inherits a stale "Stopping..." state. const stopRequested = running && stopRequestedLocal; - // Whether there is a run to show live (running or a not-yet-dismissed completed run). - // When false, the progress card + charts still render but grayed, as a preview. + // Whether there is a run to show live: running, or ANY terminal run (completed / + // stopped / error) the user has not dismissed yet. Dismissing must cover every + // terminal status, or "Train another" after a stop (and any error) would trap the + // run view with no way back to the settings. + const terminalStatuses = ["completed", "stopped", "error"]; const hasRun = Boolean( status && status.status !== "idle" && - !(status.status === "completed" && status.job_id === dismissedJobId), + !(terminalStatuses.includes(status.status) && status.job_id === dismissedJobId), ); // Notify the parent exactly once per run that produced an adapter (full completion or @@ -1226,6 +1229,22 @@ export function DiffusionTrainPanel({ {stopRequested ? "Stopping..." : "Stop training"} )} + {/* Terminal runs WITHOUT an adapter card (error, or a stop that discarded + the run) still need a way back to the settings. */} + {!running && + status && + terminalStatuses.includes(status.status) && + !completed && + !stoppedWithAdapter && ( + + )}
{(completed || stoppedWithAdapter) && ( From 44c6984f7397917bdabd2c86f7a1786b51e9bf2d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 3 Jul 2026 11:58:24 +0000 Subject: [PATCH 09/24] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../training/diffusion_training_service.py | 4 +-- studio/backend/routes/training.py | 1 - .../backend/tests/test_diffusion_training.py | 28 ++++++++++++++----- 3 files changed, 22 insertions(+), 11 deletions(-) diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index 0f57f4d4ca..fb1527b8bf 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -77,9 +77,7 @@ 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 - ) + 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] = [] diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index a25983dcc3..1a85089439 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1285,7 +1285,6 @@ async def list_diffusion_training_runs( """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 - return DiffusionTrainingRunsResponse( runs = [DiffusionTrainingRunSummary(**r) for r in list_diffusion_runs(limit = limit)] ) diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index ad43cf85f0..c9bb62335a 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -670,7 +670,9 @@ def test_run_record_no_save_stop_marks_unsaved(_isolated_runs_dir): 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}) + 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)) @@ -692,14 +694,26 @@ def test_runs_endpoints_list_and_detail(client, _isolated_runs_dir): 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]}, + "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]}, + "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" From df1ecee819173618fa33a7e583a2e76372af8cc6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 13:25:16 +0000 Subject: [PATCH 10/24] Train tab: epochs run length + stop dialog wrap fix - num_epochs on the diffusion train request and config: > 0 overrides train_steps with epochs x ceil(N / (batch x grad_accum)) optimizer steps, resolved against the dataset size in both the DiT and SDXL trainers - Train settings: run length control with a Steps / Epochs unit select - Stop dialog: flex-wrap footer so Stop and save wraps instead of clipping out of frame at narrow window widths --- .../core/training/diffusion_dit_trainer.py | 7 ++- .../core/training/diffusion_lora_trainer.py | 6 ++ .../core/training/diffusion_train_common.py | 22 +++++++ studio/backend/models/training.py | 9 +++ .../tests/test_diffusion_lora_trainer.py | 55 ++++++++++++++++ .../backend/tests/test_diffusion_training.py | 23 +++++++ studio/frontend/src/features/images/api.ts | 3 + .../images/train/diffusion-train-panel.tsx | 62 +++++++++++++++++-- 8 files changed, 180 insertions(+), 7 deletions(-) diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 5776f87652..44597ce561 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -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 @@ -48,6 +48,7 @@ from core.training.diffusion_train_common import ( _restore_perf_flags, discover_image_caption_pairs, repo_is_prequantized, + resolve_train_steps, ) # Per-family LoRA target modules (attention projections). FLUX / Qwen double-stream blocks @@ -948,6 +949,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() diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index 84bbb6bf74..09d1c5618b 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -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 @@ -58,6 +59,7 @@ from core.training.diffusion_train_common import ( # noqa: F401 _restore_perf_flags, discover_image_caption_pairs, get_trainer, + resolve_train_steps, ) @@ -301,6 +303,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 diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 56b1020cd5..05faaa5f3b 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -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 @@ -245,6 +246,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 @@ -298,6 +302,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: @@ -359,6 +365,18 @@ 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 + + def discover_image_caption_pairs( data_dir: str | os.PathLike[str], *, @@ -596,6 +614,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", diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index d7703f7141..aa7b1f29fe 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -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) diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py index ff8f3ec89e..3465408a0e 100644 --- a/studio/backend/tests/test_diffusion_lora_trainer.py +++ b/studio/backend/tests/test_diffusion_lora_trainer.py @@ -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, ) @@ -102,6 +103,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) diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index c9bb62335a..6320b7fc42 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -316,6 +316,29 @@ 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_route_start_rejects_uncontained_paths(client): # An absolute path outside the Studio dataset roots is a 400, not silently accepted. r = client.post("/api/train/diffusion/start", json = {**_BODY, "data_dir": "/etc"}) diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index 61726a02ff..eb8ecd785e 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -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; diff --git a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx index 0a9f44b839..a76e8e5209 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -210,6 +210,10 @@ export function DiffusionTrainPanel({ const [instancePrompt, setInstancePrompt] = useState(""); const [steps, setSteps] = useState(500); + // Run length is set in either steps or epochs; the trainer resolves epochs -> steps once + // the dataset size is known (num_epochs overrides train_steps on the backend). + const [durationUnit, setDurationUnit] = useState<"steps" | "epochs">("steps"); + const [epochs, setEpochs] = useState(10); const [learningRate, setLearningRate] = useState(family?.defaults.lr ?? 0.0001); const [rank, setRank] = useState(family?.defaults.rank ?? 16); const [resolution, setResolution] = useState(family?.defaults.resolution ?? 768); @@ -553,7 +557,11 @@ export function DiffusionTrainPanel({ ); return; } - if (steps < 1) return toast.error("Steps must be at least 1."); + if (durationUnit === "epochs") { + if (epochs < 1) return toast.error("Epochs must be at least 1."); + } else if (steps < 1) { + return toast.error("Steps must be at least 1."); + } if (rank < 1) return toast.error("LoRA rank must be at least 1."); if (resolution < 64 || resolution % 8 !== 0) { return toast.error("Resolution must be a multiple of 8 and at least 64."); @@ -577,7 +585,10 @@ export function DiffusionTrainPanel({ output_dir: outputDir.trim(), instance_prompt: instancePrompt.trim() || undefined, resolution, - train_steps: steps, + // Epochs mode overrides train_steps on the backend, so send num_epochs and omit + // train_steps (the backend default is unused when num_epochs > 0). + train_steps: durationUnit === "epochs" ? undefined : steps, + num_epochs: durationUnit === "epochs" ? epochs : undefined, learning_rate: learningRate, train_batch_size: batchSize, gradient_accumulation_steps: gradAccum, @@ -610,6 +621,8 @@ export function DiffusionTrainPanel({ instancePrompt, resolution, steps, + durationUnit, + epochs, learningRate, batchSize, gradAccum, @@ -690,6 +703,40 @@ export function DiffusionTrainPanel({
); + // Run length: a number paired with a compact unit select (Steps / Epochs). Epochs mode + // trains for that many full passes over the dataset; the backend resolves it to steps. + const durationField = ( +
+ +
+ { + settingsDirty.current = true; + const n = Number(e.target.value) || 1; + if (durationUnit === "epochs") setEpochs(n); + else setSteps(n); + }} + className="h-8 min-w-0 flex-1 text-xs" + /> + +
+
+ ); + const precisionLabel = (m: "nf4" | "bf16" | "int8" | "fp8" | "auto"): string => { if (m === "auto") return "Auto (recommended)"; if (m === "nf4") return "nf4 (4-bit QLoRA, lowest VRAM)"; @@ -704,7 +751,7 @@ export function DiffusionTrainPanel({ const trainingSettings = (
- {numberField("Steps", steps, setSteps, 1)} + {durationField} {numberField("LoRA rank", rank, setRank, 1)} {numberField("Resolution", resolution, setResolution, 512, { min: 64, step: 64 })} {numberField("Batch", batchSize, setBatchSize, 1)} @@ -1292,9 +1339,12 @@ export function DiffusionTrainPanel({ finishes first. - {/* items-center + a real label on the destructive action: a bare "Stop" rendered - as a stubby pill between two wide ones and read as misaligned. */} - + {/* flex-wrap keeps all three buttons visible when the sm:flex-row row is wider than + the dialog at narrow widths (down to ~480px); it wraps instead of clipping the + last button past the right edge. items-center + a real label on the destructive + action: a bare "Stop" rendered as a stubby pill between two wide ones and read + as misaligned. */} + Continue training void onStop(false)}> Stop without saving From d927bf155c2bd8d04fd948e1f3a900bae9873724 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 4 Jul 2026 01:06:42 +0000 Subject: [PATCH 11/24] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/sd_cpp_backend.py | 5 ++++- studio/backend/tests/test_sd_cpp_backend.py | 14 +++++++++++--- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 7632551661..d8cb21d409 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -893,7 +893,10 @@ class SdCppDiffusionBackend: lora_stage = Path(server_lora_dir) / f"gen_{os.urandom(6).hex()}" materialized = diffusion_lora.materialize_native_dir(lora_resolved, lora_stage) lora_payload = [ - {"path": f"{lora_stage.name}/{Path(m.path).name}", "multiplier": float(m.weight)} + { + "path": f"{lora_stage.name}/{Path(m.path).name}", + "multiplier": float(m.weight), + } for m in materialized ] try: diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py index 0b52ea6e14..be2fc22a4b 100644 --- a/studio/backend/tests/test_sd_cpp_backend.py +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -691,7 +691,11 @@ def _fake_materialize(resolved, dest): return out -def _patch_lora(monkeypatch, resolved, supported = True): +def _patch_lora( + monkeypatch, + resolved, + supported = True, +): from core.inference import diffusion_lora as dl monkeypatch.setattr(dl, "supports_lora", lambda **k: supported) @@ -706,7 +710,9 @@ def test_generate_oneshot_applies_loras_via_prompt_tags(monkeypatch): eng = _FakeEngine() b = _loaded_backend(engine = eng) # mode = "oneshot" - _patch_lora(monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.8)]) + _patch_lora( + monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.8)] + ) b.generate(prompt = "a fox", steps = 4, seed = 1, loras = [("id1", 0.8)]) _, params, _, _ = eng.calls[0] assert params.lora_dir is not None and params.lora_apply_mode == "auto" @@ -725,7 +731,9 @@ def test_generate_server_stages_loras_and_sends_structured_field(monkeypatch, tm servers: list = [] _run_server_load(monkeypatch, b, servers) servers[0].lora_dir = str(tmp_path) - _patch_lora(monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.7)]) + _patch_lora( + monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.7)] + ) b.generate(prompt = "x", steps = 4, seed = 1, batch_size = 1, loras = [("id1", 0.7)]) payload = servers[0].payloads[0] assert "lora" in payload and len(payload["lora"]) == 1 From 23c6457e62a4ef026fbaff24714df469bbc5187f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 01:27:40 +0000 Subject: [PATCH 12/24] Address review: Train panel precision and notification edge cases - Reset mixed precision to bf16 when the family changes to a DiT: an fp16/no value left over from SDXL rode along in the DiT start payload and the backend rejected it (dense base precisions require bf16 compute). - Gate the dense base precisions behind the selected base: a bnb-4bit repo disables bf16/int8/fp8 with a hint, and a dense selection auto-flips to auto so the run does not fail at the validator. - Re-arm the run-completion notification in onStart, so a second run notifies even when its running phase is never observed by the poll. --- .../images/train/diffusion-train-panel.tsx | 64 ++++++++++++++++++- 1 file changed, 62 insertions(+), 2 deletions(-) diff --git a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx index a76e8e5209..813a6c1fd2 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -92,6 +92,21 @@ const FAMILY_PRESETS: FamilyPreset[] = [ const CUSTOM_BASE = "__custom__"; const UPLOAD_DATASET = "__upload__"; +// The dense DiT base precisions: they load a dense (bf16) base and quantise/cast it, so the +// backend rejects them for an already-quantised bnb-4bit repo. "nf4"/"auto" stay valid. +const DENSE_PRECISIONS = new Set(["bf16", "int8", "fp8"]); +// Mirror the backend's repo_is_prequantized heuristic: a repo whose name marks a +// bitsandbytes 4-bit build already ships a quantised transformer and cannot serve the dense +// base precisions. Kept in sync with diffusion_train_common.repo_is_prequantized. +function repoIsPrequantized(baseModel: string): boolean { + const name = baseModel.toLowerCase(); + return ( + name.includes("bnb-4bit") || + name.includes("-4bit") || + name.includes("int4") || + name.includes("nf4") + ); +} // Dataset-select option value prefix for a not-yet-imported example; picking it imports. const EXAMPLE_PREFIX = "example:"; const DATASET_FILE_ACCEPT = ".png,.jpg,.jpeg,.webp,.bmp,.txt,.caption,.jsonl"; @@ -372,6 +387,16 @@ export function DiffusionTrainPanel({ } }, [family, loadedBaseRepo, reportedFamily?.recommended_precision]); + // mixed_precision is an SDXL-only lever (its UI control is hidden for DiT families). A + // dense DiT base precision (bf16/int8/fp8) requires bf16 compute, and every DiT family + // trains in bf16, so reset precision to bf16 when the family changes to a DiT. Without + // this, an fp16/no value left over from SDXL rides along in the DiT start payload and the + // backend rejects it (dense modes need mixed_precision=bf16). Kept in its own effect so it + // does not re-trigger the base/settings reseed above. + useEffect(() => { + if (isDiT) setPrecision("bf16"); + }, [isDiT]); + // The base actually used everywhere (request, deploy, select value). baseChoice can // briefly hold another family's repo between a family switch and the reseed effect // (or if that effect is skipped); a raw

- How the learning rate evolves over the run (shown live in the LR chart). + How the learning rate evolves over the run.

{lrScheduler !== "constant" &&