diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index 4963714924..a36c345774 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -319,3 +319,41 @@ export async function stopDiffusionTraining(): Promise<{ status: string }> { 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 index 85fe0dd99b..25b260a03e 100644 --- a/studio/frontend/src/features/images/diffusion-train-dialog.tsx +++ b/studio/frontend/src/features/images/diffusion-train-dialog.tsx @@ -1,7 +1,7 @@ // 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 { useCallback, useEffect, useRef, useState } from "react"; import { Button } from "@/components/ui/button"; import { @@ -18,17 +18,32 @@ 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"; -const DEFAULT_SDXL_BASE = "stabilityai/stable-diffusion-xl-base-1.0"; +// 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"; -// 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. +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, @@ -41,10 +56,16 @@ export function DiffusionTrainDialog({ // Called once when a run finishes so the page can rescan the LoRA picker. onTrainingComplete?: () => void; }) { - const [baseModel, setBaseModel] = useState(defaultBaseModel || DEFAULT_SDXL_BASE); - const [dataDir, setDataDir] = useState(""); + 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); @@ -54,13 +75,33 @@ export function DiffusionTrainDialog({ 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. + // 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) setBaseModel(defaultBaseModel || DEFAULT_SDXL_BASE); - }, [open, defaultBaseModel]); + 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 { @@ -79,6 +120,7 @@ export function DiffusionTrainDialog({ }, [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)) @@ -97,9 +139,56 @@ export function DiffusionTrainDialog({ } }, [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 () => { - if (!baseModel.trim() || !dataDir.trim() || !outputDir.trim()) { - toast.error("Base model, dataset folder, and output folder are required."); + 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 @@ -114,8 +203,8 @@ export function DiffusionTrainDialog({ setStarting(true); try { await startDiffusionTraining({ - base_model: baseModel.trim(), - data_dir: dataDir.trim(), + base_model: baseModel, + data_dir: dataset, output_dir: outputDir.trim(), instance_prompt: instancePrompt.trim() || undefined, resolution, @@ -136,8 +225,10 @@ export function DiffusionTrainDialog({ setStarting(false); } }, [ - baseModel, - dataDir, + baseChoice, + customBase, + dataset, + selectedDataset, outputDir, instancePrompt, resolution, @@ -163,122 +254,211 @@ export function DiffusionTrainDialog({ - Train a LoRA (SDXL) + Train an SDXL LoRA - 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. Folders resolve inside the - Studio home: datasets under its datasets folder, the adapter under its outputs - folder (the exact save path is shown after the run). + 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. */}
- - setBaseModel(e.target.value)} - className="h-8 text-xs" - /> + + + {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. */}
- - setDataDir(e.target.value)} - className="h-8 text-xs" - /> + + + {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" />
-
-
- - 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" - /> -
-
- - -
-
+ + {/* 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" && (
@@ -292,8 +472,10 @@ export function DiffusionTrainDialog({
- {status.message} - {status.loss != null && <> · loss {status.loss.toFixed(4)}} + {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}
)} @@ -307,8 +489,17 @@ export function DiffusionTrainDialog({ + ) : completed ? ( + <> + + + ) : ( - )} diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index fcad1bb8ed..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, @@ -1917,8 +1918,9 @@ export function ImagesPage({ active = true }: { active?: boolean }) { size="sm" className="h-[34px]" onClick={() => setTrainOpen(true)} - title="Train a LoRA adapter (SDXL)" + title="Teach SDXL your own style or subject from a folder of images" > + Train LoRA {/* Single fixed toggle for the right-docked Advanced panel (mirrors Chat's settings