From 96575cebf5ec77a08e94d23bde447fa48a577262 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 01:14:28 +0000 Subject: [PATCH] Images Train LoRA dialog: token, validation, precision, base-repo prefill, gating, refresh Nine review findings on the SDXL training dialog: - Forward the saved Hub token so a gated/private SDXL base can be trained (the image load flow already sends it). - Re-seed the base-model field from the current default each time the dialog opens; the keep-alive dialog otherwise kept its mount-time default after a model loaded. - Prefill from base_repo (the diffusers pipeline) rather than repo_id, which for a GGUF/single-file SDXL load is the checkpoint path from_pretrained can't open. - Add client-side validation of steps/rank/resolution/batch/learning-rate before the request. - Expose a precision selector (bf16/fp16/fp32) so non-bf16 GPUs can train from the UI, not only the API. - Gate the dialog on the active Images route (active && trainOpen) so switching tabs closes it and stops its polling. - Rescan the LoRA picker when a run completes, so a freshly-trained adapter appears without a model reload. - Cap the dialog height and scroll the body so the Start/Stop footer stays reachable on short viewports. - Correct the copy to not over-promise picker auto-discovery. Freeing the resident Images pipeline before training is handled backend-side in the diffusion training start route. --- studio/frontend/src/features/images/api.ts | 4 + .../images/diffusion-train-dialog.tsx | 99 ++++++++++++++++--- .../src/features/images/images-page.tsx | 17 +++- 3 files changed, 101 insertions(+), 19 deletions(-) diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index 26e9e25bed..efc57aa862 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -269,10 +269,14 @@ export interface DiffusionTrainingStartRequest { gradient_accumulation_steps?: number; lora_rank?: number; lora_alpha?: number | null; + lora_target_modules?: string[]; + max_grad_norm?: number; seed?: number; mixed_precision?: "bf16" | "fp16" | "no"; gradient_checkpointing?: boolean; lr_scheduler?: string; + // Forwarded to StableDiffusionXLPipeline.from_pretrained for a gated/private base repo. + hf_token?: string | null; } // A snapshot of the current diffusion training job (GET /api/train/diffusion/status). diff --git a/studio/frontend/src/features/images/diffusion-train-dialog.tsx b/studio/frontend/src/features/images/diffusion-train-dialog.tsx index a365908dae..cf8fdb4b90 100644 --- a/studio/frontend/src/features/images/diffusion-train-dialog.tsx +++ b/studio/frontend/src/features/images/diffusion-train-dialog.tsx @@ -14,6 +14,7 @@ import { } from "@/components/ui/dialog"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; +import { getHfToken, hfApiToken } from "@/features/hub/stores/hf-token-store"; import { toast } from "@/lib/toast"; import { @@ -23,6 +24,8 @@ import { stopDiffusionTraining, } from "./api"; +const DEFAULT_SDXL_BASE = "stabilityai/stable-diffusion-xl-base-1.0"; + // A self-contained "Train a LoRA" dialog for the diffusion (SDXL) trainer. It posts to // /api/train/diffusion/start and polls /status while open, so it never blocks the page and // works whether or not a model is loaded for generation. Only SDXL is trainable today. @@ -30,12 +33,15 @@ export function DiffusionTrainDialog({ open, onOpenChange, defaultBaseModel, + onTrainingComplete, }: { open: boolean; onOpenChange: (open: boolean) => void; defaultBaseModel?: string; + // Called once when a run finishes so the page can rescan the LoRA picker. + onTrainingComplete?: () => void; }) { - const [baseModel, setBaseModel] = useState(defaultBaseModel || "stabilityai/stable-diffusion-xl-base-1.0"); + const [baseModel, setBaseModel] = useState(defaultBaseModel || DEFAULT_SDXL_BASE); const [dataDir, setDataDir] = useState(""); const [outputDir, setOutputDir] = useState(""); const [instancePrompt, setInstancePrompt] = useState(""); @@ -44,9 +50,18 @@ export function DiffusionTrainDialog({ const [rank, setRank] = useState(16); const [resolution, setResolution] = useState(1024); const [batchSize, setBatchSize] = useState(1); + const [precision, setPrecision] = useState<"bf16" | "fp16" | "no">("bf16"); const [starting, setStarting] = useState(false); const [status, setStatus] = useState(null); + // The dialog stays mounted (ImagesPage is keep-alive), so the initial state seed does not + // reflect a base model loaded AFTER mount. Re-seed the base-model field from the current + // default each time the dialog opens, so "Train LoRA" after loading an SDXL checkpoint + // starts from that checkpoint's diffusers repo, not the hard-coded default. + useEffect(() => { + if (open) setBaseModel(defaultBaseModel || DEFAULT_SDXL_BASE); + }, [open, defaultBaseModel]); + const poll = useCallback(async () => { try { setStatus(await getDiffusionTrainingStatus()); @@ -69,11 +84,33 @@ export function DiffusionTrainDialog({ ? Math.min(100, Math.round((status.step / status.total_steps) * 100)) : 0; + // Notify the parent exactly once when a run reaches "completed", so it can rescan the + // LoRA picker (a LoRA trained while a model is loaded is otherwise invisible until a + // model swap re-runs the discovery effect). + const [notifiedComplete, setNotifiedComplete] = useState(false); + useEffect(() => { + if (status?.status === "completed" && !notifiedComplete) { + setNotifiedComplete(true); + onTrainingComplete?.(); + } else if (status?.status === "running" && notifiedComplete) { + setNotifiedComplete(false); // arm again for the next run + } + }, [status?.status, notifiedComplete, onTrainingComplete]); + const onStart = useCallback(async () => { if (!baseModel.trim() || !dataDir.trim() || !outputDir.trim()) { toast.error("Base model, dataset folder, and output folder are required."); return; } + // Mirror the backend's numeric validation so obvious mistakes are caught before the + // request (the backend returns 400 for these; catching here gives a clearer message). + 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."); + } + if (batchSize < 1) return toast.error("Batch size must be at least 1."); + if (learningRate <= 0) return toast.error("Learning rate must be greater than 0."); setStarting(true); try { await startDiffusionTraining({ @@ -86,6 +123,10 @@ export function DiffusionTrainDialog({ learning_rate: learningRate, train_batch_size: batchSize, lora_rank: rank, + mixed_precision: precision, + // Forward the saved Hub token so a gated/private SDXL base can be trained (the + // image load flow already sends it, so a model you can load, you can also train). + hf_token: hfApiToken(getHfToken()) || undefined, }); toast.success("Training started"); void poll(); @@ -94,7 +135,19 @@ export function DiffusionTrainDialog({ } finally { setStarting(false); } - }, [baseModel, dataDir, outputDir, instancePrompt, resolution, steps, learningRate, batchSize, rank, poll]); + }, [ + baseModel, + dataDir, + outputDir, + instancePrompt, + resolution, + steps, + learningRate, + batchSize, + rank, + precision, + poll, + ]); const onStop = useCallback(async () => { try { @@ -108,17 +161,17 @@ export function DiffusionTrainDialog({ return ( - + Train a LoRA (SDXL) Fine-tune an SDXL LoRA on a folder of images. Captions come from a metadata.jsonl, - per-image .txt sidecars, or the instance prompt below. The adapter is written to the - output folder and can be loaded from the LoRAs picker. + per-image .txt sidecars, or the instance prompt below. The adapter is saved to the + output folder shown after the run. -
+
-
- - setLearningRate(Number(e.target.value) || 0.0001)} - className="h-8 text-xs" - /> +
+
+ + setLearningRate(Number(e.target.value) || 0.0001)} + className="h-8 text-xs" + /> +
+
+ + +
{status && status.status !== "idle" && ( diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 4df7aef66d..785ea68dbc 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -910,6 +910,9 @@ export function ImagesPage({ active = true }: { active?: boolean }) { const [availableLoras, setAvailableLoras] = useState([]); // "Train a LoRA" dialog (SDXL). Independent of the loaded generation model. const [trainOpen, setTrainOpen] = useState(false); + // Bumped when a training run completes, to force the LoRA discovery effect to rescan so + // a freshly-trained adapter appears in the picker without a model reload. + const [loraRefreshKey, setLoraRefreshKey] = useState(0); // ControlNet for the next generation: the chosen model id, a control image (data URL), // how to derive the control map, and the conditioning strength. Available models refresh // per loaded family; applied at generate time only when a model + control image are set. @@ -1011,7 +1014,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { return () => { cancelled = true; }; - }, [loraCapable, status?.family]); + }, [loraCapable, status?.family, loraRefreshKey]); // Refresh the ControlNet picker's options when the loaded model (family) changes, and clear // a stale selection the new model can't use so an incompatible ControlNet is never sent. @@ -1768,9 +1771,17 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
setLoraRefreshKey((k) => k + 1)} /> {/* ── Controls rail + preview canvas. Padding mirrors the other tabs