From 906f541ac91f075a456cb0146c0b7b6a080f57a3 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 09:49:23 +0000 Subject: [PATCH] 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 + + + +
); }