From 556b4cc346b8effe41f7a6b79d22eea3ca4c60ca Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 15:18:50 +0000 Subject: [PATCH] Images: add a Train LoRA (SDXL) dialog Surface the diffusion training API in the Images page. A "Train LoRA" button in the top bar opens a self-contained dialog to fine-tune an SDXL LoRA on a folder of images: pick the base model, dataset folder, output folder, an optional instance prompt, and the core hyperparameters (steps, rank, resolution, batch, learning rate), then Start. The dialog polls the training status while open and shows a progress bar, step count, live loss, and the saved adapter path, with a Stop button for a clean stop. The dialog is independent of the loaded generation model (training runs in its own subprocess), and prefills the base model with the loaded checkpoint when it is SDXL, else the SDXL base. api.ts gains startDiffusionTraining / stopDiffusionTraining / getDiffusionTrainingStatus plus their types, matching the /api/train/diffusion routes. --- studio/frontend/src/features/images/api.ts | 59 ++++ .../images/diffusion-train-dialog.tsx | 251 ++++++++++++++++++ .../src/features/images/images-page.tsx | 56 ++-- 3 files changed, 349 insertions(+), 17 deletions(-) create mode 100644 studio/frontend/src/features/images/diffusion-train-dialog.tsx diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index e33da62113..26e9e25bed 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -254,3 +254,62 @@ export async function fetchGalleryObjectUrl(url: string): Promise { if (!res.ok) throw new Error(await readFastApiError(res)); return URL.createObjectURL(await res.blob()); } + +// ── Diffusion (SDXL) LoRA training ──────────────────────────────────────────── +// Mirrors DiffusionTrainingStartRequest on the backend; only the paths are required. +export interface DiffusionTrainingStartRequest { + base_model: string; + data_dir: string; + output_dir: string; + instance_prompt?: string | null; + resolution?: number; + train_steps?: number; + learning_rate?: number; + train_batch_size?: number; + gradient_accumulation_steps?: number; + lora_rank?: number; + lora_alpha?: number | null; + seed?: number; + mixed_precision?: "bf16" | "fp16" | "no"; + gradient_checkpointing?: boolean; + lr_scheduler?: string; +} + +// A snapshot of the current diffusion training job (GET /api/train/diffusion/status). +export interface DiffusionTrainingStatus { + active: boolean; + job_id: string | null; + status: string; + message: string; + step: number; + total_steps: number; + loss: number | null; + avg_loss: number | null; + learning_rate: number | null; + num_images: number | null; + in_model_load: boolean; + output_dir: string | null; + lora_path: string | null; + started_at: number | null; + updated_at: number | null; +} + +export async function startDiffusionTraining( + body: DiffusionTrainingStartRequest, +): Promise<{ job_id: string; status: string }> { + return parseJson( + await authFetch("/api/train/diffusion/start", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(body), + }), + ); +} + +export async function stopDiffusionTraining(): Promise<{ status: string }> { + return parseJson(await authFetch("/api/train/diffusion/stop", { method: "POST" })); +} + +export async function getDiffusionTrainingStatus(): Promise { + return parseJson(await authFetch("/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 new file mode 100644 index 0000000000..a365908dae --- /dev/null +++ b/studio/frontend/src/features/images/diffusion-train-dialog.tsx @@ -0,0 +1,251 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { useCallback, useEffect, useState } from "react"; + +import { Button } from "@/components/ui/button"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { toast } from "@/lib/toast"; + +import { + type DiffusionTrainingStatus, + getDiffusionTrainingStatus, + startDiffusionTraining, + stopDiffusionTraining, +} from "./api"; + +// 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. +export function DiffusionTrainDialog({ + open, + onOpenChange, + defaultBaseModel, +}: { + open: boolean; + onOpenChange: (open: boolean) => void; + defaultBaseModel?: string; +}) { + const [baseModel, setBaseModel] = useState(defaultBaseModel || "stabilityai/stable-diffusion-xl-base-1.0"); + const [dataDir, setDataDir] = useState(""); + const [outputDir, setOutputDir] = useState(""); + const [instancePrompt, setInstancePrompt] = useState(""); + const [steps, setSteps] = useState(500); + const [learningRate, setLearningRate] = useState(0.0001); + const [rank, setRank] = useState(16); + const [resolution, setResolution] = useState(1024); + const [batchSize, setBatchSize] = useState(1); + const [starting, setStarting] = useState(false); + const [status, setStatus] = useState(null); + + const poll = useCallback(async () => { + try { + setStatus(await getDiffusionTrainingStatus()); + } catch { + // Best-effort; a failed poll should not surface an error while the dialog is open. + } + }, []); + + // Poll status only while the dialog is open. + useEffect(() => { + if (!open) return; + void poll(); + const id = window.setInterval(() => void poll(), 1500); + return () => window.clearInterval(id); + }, [open, poll]); + + const active = Boolean(status?.active) || status?.status === "running"; + const pct = + status && status.total_steps > 0 + ? Math.min(100, Math.round((status.step / status.total_steps) * 100)) + : 0; + + const onStart = useCallback(async () => { + if (!baseModel.trim() || !dataDir.trim() || !outputDir.trim()) { + toast.error("Base model, dataset folder, and output folder are required."); + return; + } + setStarting(true); + try { + await startDiffusionTraining({ + base_model: baseModel.trim(), + data_dir: dataDir.trim(), + output_dir: outputDir.trim(), + instance_prompt: instancePrompt.trim() || undefined, + resolution, + train_steps: steps, + learning_rate: learningRate, + train_batch_size: batchSize, + lora_rank: rank, + }); + toast.success("Training started"); + void poll(); + } catch (e) { + toast.error(e instanceof Error ? e.message : "Failed to start training"); + } finally { + setStarting(false); + } + }, [baseModel, dataDir, outputDir, instancePrompt, resolution, steps, learningRate, batchSize, rank, 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]); + + 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. + + + +
+
+ + setBaseModel(e.target.value)} + className="h-8 text-xs" + /> +
+
+ + setDataDir(e.target.value)} + className="h-8 text-xs" + /> +
+
+ + setOutputDir(e.target.value)} + className="h-8 text-xs" + /> +
+
+ + setInstancePrompt(e.target.value)} + className="h-8 text-xs" + /> +
+
+
+ + setSteps(Number(e.target.value) || 1)} + className="h-8 text-xs" + /> +
+
+ + setRank(Number(e.target.value) || 1)} + className="h-8 text-xs" + /> +
+
+ + setResolution(Number(e.target.value) || 1024)} + className="h-8 text-xs" + /> +
+
+ + setBatchSize(Number(e.target.value) || 1)} + className="h-8 text-xs" + /> +
+
+
+ + setLearningRate(Number(e.target.value) || 0.0001)} + className="h-8 text-xs" + /> +
+ + {status && status.status !== "idle" && ( +
+
+ {status.status} + + {status.total_steps > 0 ? `${status.step}/${status.total_steps}` : ""} + +
+
+
+
+
+ {status.message} + {status.loss != null && <> · loss {status.loss.toFixed(4)}} + {status.lora_path && ( +
Saved: {status.lora_path}
+ )} +
+
+ )} +
+ + + {active ? ( + + ) : ( + + )} + + +
+ ); +} diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index eacfea5c15..4df7aef66d 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -66,6 +66,7 @@ import { loadDiffusionModel, unloadDiffusionModel, } from "./api"; +import { DiffusionTrainDialog } from "./diffusion-train-dialog"; // Curated diffusion GGUFs the picker recommends. The backend resolves each one's // pipeline + base diffusers repo from its repo id, so the rail just lists them; @@ -907,6 +908,8 @@ export function ImagesPage({ active = true }: { active?: boolean }) { // offers. Applied at generate time; available adapters are refreshed per loaded family. const [loras, setLoras] = useState([]); const [availableLoras, setAvailableLoras] = useState([]); + // "Train a LoRA" dialog (SDXL). Independent of the loaded generation model. + const [trainOpen, setTrainOpen] = useState(false); // 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. @@ -1732,24 +1735,43 @@ export function ImagesPage({ active = true }: { active?: boolean }) { open={active && selectorOpen} onOpenChange={(o) => setSelectorOpen(active && o)} /> - {/* 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. */} - +
+ {/* Train a LoRA (SDXL): opens a self-contained dialog; available regardless of + whether a generation model is loaded. */} + + {/* 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. */} + +
+ {/* ── Controls rail + preview canvas. Padding mirrors the other tabs (Export, Data Recipes): px-5 / sm:px-9, with a roomy bottom. ── */}