diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py index e238fdaff7..fd17b29500 100644 --- a/studio/backend/tests/test_diffusion_lora_trainer.py +++ b/studio/backend/tests/test_diffusion_lora_trainer.py @@ -208,9 +208,7 @@ def test_config_rejects_known_non_sdxl_base_models(): "z-image-turbo-Q4_K_M.gguf", ): with pytest.raises(ValueError, match = "SDXL"): - DiffusionLoraConfig( - base_model = bad, data_dir = "d", output_dir = "o" - ).normalized() + DiffusionLoraConfig(base_model = bad, data_dir = "d", output_dir = "o").normalized() def test_config_accepts_sdxl_and_unknown_base_models(): diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index ad38261705..a36c345774 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -256,3 +256,104 @@ 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; + 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). +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")); +} + +// One image-dataset folder under the Studio datasets root (GET /api/train/diffusion/info). +export interface DiffusionDatasetSummary { + name: string; + path: string; + image_count: number; + caption_count: number; +} + +// Where diffusion training reads/writes on this Studio, plus usable dataset folders. +export interface DiffusionTrainingInfo { + datasets_root: string; + outputs_root: string; + datasets: DiffusionDatasetSummary[]; +} + +export async function getDiffusionTrainingInfo(): Promise { + return parseJson(await authFetch("/api/train/diffusion/info")); +} + +export interface DiffusionDatasetUploadResult extends DiffusionDatasetSummary { + uploaded: number; +} + +/** Upload images (+ optional caption .txt / metadata.jsonl) into a named dataset folder. + * Repeat uploads into the same name accumulate; the returned name is a valid data_dir + * for startDiffusionTraining. */ +export async function uploadDiffusionDataset( + name: string, + files: File[], +): Promise { + const form = new FormData(); + form.append("name", name); + for (const f of files) form.append("files", f); + return parseJson( + await authFetch("/api/train/diffusion/dataset", { method: "POST", body: form }), + ); +} 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..25b260a03e --- /dev/null +++ b/studio/frontend/src/features/images/diffusion-train-dialog.tsx @@ -0,0 +1,510 @@ +// 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, useRef, 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 { getHfToken, hfApiToken } from "@/features/hub/stores/hf-token-store"; +import { toast } from "@/lib/toast"; + +import { + type DiffusionTrainingInfo, + type DiffusionTrainingStatus, + getDiffusionTrainingInfo, + getDiffusionTrainingStatus, + startDiffusionTraining, + stopDiffusionTraining, + uploadDiffusionDataset, +} from "./api"; + +// The two official SDXL bases the backend allowlists for non-GGUF loads. Everything the +// dropdown offers is trainable; "custom" is the escape hatch for local SDXL checkpoints. +const SDXL_BASES: Array<{ id: string; label: string }> = [ + { id: "stabilityai/stable-diffusion-xl-base-1.0", label: "SDXL Base 1.0 (best quality)" }, + { id: "stabilityai/sdxl-turbo", label: "SDXL Turbo (fast, good for quick tests)" }, +]; +const CUSTOM_BASE = "__custom__"; +const UPLOAD_DATASET = "__upload__"; +const DATASET_FILE_ACCEPT = ".png,.jpg,.jpeg,.webp,.bmp,.txt,.caption,.jsonl"; + +const selectClass = + "h-8 w-full rounded-md border border-input bg-background px-2 text-xs"; + +// A self-contained "Train an SDXL LoRA" dialog. 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; the backend refuses +// known non-SDXL picks instantly, and the base-model dropdown keeps users on safe picks. +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 [baseChoice, setBaseChoice] = useState(SDXL_BASES[0].id); + const [customBase, setCustomBase] = useState(""); + const [info, setInfo] = useState(null); + const [dataset, setDataset] = useState(UPLOAD_DATASET); + const [uploadName, setUploadName] = useState("my-images"); + const [uploading, setUploading] = useState(false); + const fileInputRef = useRef(null); + const [outputDir, setOutputDir] = useState(""); + const [instancePrompt, setInstancePrompt] = useState(""); + const [showAdvanced, setShowAdvanced] = useState(false); + 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 [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 seed per-open state here: + // the base-model choice from the currently loaded SDXL pipeline (when there is one), + // and the dataset list from the backend. + const refreshInfo = useCallback(async (): Promise => { + try { + const i = await getDiffusionTrainingInfo(); + setInfo(i); + return i; + } catch { + return null; // older backends: keep the upload-only flow usable + } + }, []); + + useEffect(() => { + if (!open) return; + if (defaultBaseModel) { + setBaseChoice(defaultBaseModel); + } + void refreshInfo().then((i) => { + // Preselect the only dataset, or the freshest-looking state: with no datasets + // yet, the picker sits on "Upload new images". + setDataset((cur) => { + if (cur !== UPLOAD_DATASET && i?.datasets.some((d) => d.name === cur)) return cur; + return i && i.datasets.length > 0 ? i.datasets[0].name : UPLOAD_DATASET; + }); + }); + }, [open, defaultBaseModel, refreshInfo]); + + 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 completed = status?.status === "completed"; + const pct = + status && status.total_steps > 0 + ? 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 selectedDataset = + dataset !== UPLOAD_DATASET ? info?.datasets.find((d) => d.name === dataset) : undefined; + + const onUpload = useCallback(async () => { + const files = Array.from(fileInputRef.current?.files ?? []); + if (files.length === 0) { + toast.error("Choose the images to upload first."); + return; + } + const name = uploadName.trim(); + if (!name) { + toast.error("Give the dataset a folder name, e.g. my-style-photos."); + return; + } + setUploading(true); + try { + const res = await uploadDiffusionDataset(name, files); + toast.success( + `Uploaded ${res.uploaded} file${res.uploaded === 1 ? "" : "s"} - ` + + `"${res.name}" now has ${res.image_count} images`, + ); + if (fileInputRef.current) fileInputRef.current.value = ""; + await refreshInfo(); + setDataset(res.name); + } catch (e) { + toast.error(e instanceof Error ? e.message : "Upload failed"); + } finally { + setUploading(false); + } + }, [uploadName, refreshInfo]); + + const onStart = useCallback(async () => { + const baseModel = (baseChoice === CUSTOM_BASE ? customBase : baseChoice).trim(); + if (!baseModel) { + toast.error("Pick a base model (or fill in the custom repo/path)."); + return; + } + if (dataset === UPLOAD_DATASET) { + toast.error("Upload your training images first (or pick an existing dataset)."); + return; + } + if (!outputDir.trim()) { + toast.error("Name the adapter (this becomes its folder under Studio outputs)."); + return; + } + if (selectedDataset && selectedDataset.caption_count === 0 && !instancePrompt.trim()) { + toast.error( + "These images have no captions - add a trigger prompt so the trainer knows " + + "what to learn (it becomes the caption for every image).", + ); + 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({ + base_model: baseModel, + data_dir: dataset, + output_dir: outputDir.trim(), + instance_prompt: instancePrompt.trim() || undefined, + resolution, + train_steps: steps, + 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(); + } catch (e) { + toast.error(e instanceof Error ? e.message : "Failed to start training"); + } finally { + setStarting(false); + } + }, [ + baseChoice, + customBase, + dataset, + selectedDataset, + outputDir, + instancePrompt, + resolution, + steps, + learningRate, + batchSize, + rank, + precision, + 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 an SDXL LoRA + + Teach SDXL a style, character, or subject from your own images. The finished + adapter shows up in this page's LoRA picker. Only SDXL can be trained for + now - FLUX, Qwen-Image and Z-Image load LoRAs but can't train them yet. + + + +
+ {/* 1. Base model: a constrained dropdown instead of free text, so the "SDXL + only" rule is embodied by the control. Custom stays available for local + SDXL checkpoints; a known non-SDXL pick is refused instantly by the API. */} +
+ + + {baseChoice === CUSTOM_BASE && ( + setCustomBase(e.target.value)} + className="h-8 text-xs" + /> + )} +
+ + {/* 2. Training images: pick an existing dataset folder or upload straight from + the browser - no shell access or Studio-home knowledge needed. */} +
+ + + {dataset === UPLOAD_DATASET && ( +
+ setUploadName(e.target.value)} + className="h-8 text-xs" + aria-label="New dataset name" + /> +
+ + +
+

+ 10-50 images work well. Optional captions: a .txt per image (same + filename) or a metadata.jsonl; without them the trigger prompt below + captions every image. You can upload more into the same name later. +

+
+ )} + {selectedDataset && selectedDataset.caption_count === 0 && ( +

+ No caption files in this dataset - the trigger prompt below will be used + as the caption for every image. +

+ )} +
+ + {/* 3. What to call the result + how to trigger it. */} +
+ + setOutputDir(e.target.value)} + className="h-8 text-xs" + /> +
+
+ + setInstancePrompt(e.target.value)} + className="h-8 text-xs" + /> +
+ + {/* 4. Hyperparameters, collapsed: the defaults suit a first run, and hiding + them keeps the primary flow at three decisions. */} + + {showAdvanced && ( + <> +
+
+ + 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}` : ""} + +
+
+
+
+
+ {completed + ? "Adapter ready - find it in the LoRA picker on this page." + : status.message} + {status.loss != null && !completed && <> · loss {status.loss.toFixed(4)}} + {status.lora_path && ( +
Saved: {status.lora_path}
+ )} +
+
+ )} +
+ + + {active ? ( + + ) : completed ? ( + <> + + + + ) : ( + + )} + + +
+ ); +} diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 4ce16eceac..54775145ce 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -3,6 +3,7 @@ import { useCallback, useEffect, useMemo, useRef, useState, type ReactNode } from "react"; import { + AiMagicIcon, ArrowLeftRightIcon, ArrowReloadHorizontalIcon, Delete02Icon, @@ -66,6 +67,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; @@ -946,6 +948,11 @@ 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); + // 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. @@ -1074,7 +1081,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. @@ -1902,24 +1909,52 @@ 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. */} + +
+ setLoraRefreshKey((k) => k + 1)} + /> {/* ── Controls rail + preview canvas. Padding mirrors the other tabs (Export, Data Recipes): px-5 / sm:px-9, with a roomy bottom. ── */}