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. ── */}