diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index a36c345774..911a2ff426 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -257,10 +257,14 @@ export async function fetchGalleryObjectUrl(url: string): Promise { return URL.createObjectURL(await res.blob()); } -// ── Diffusion (SDXL) LoRA training ──────────────────────────────────────────── +// ── Diffusion LoRA training ─────────────────────────────────────────────────── // Mirrors DiffusionTrainingStartRequest on the backend; only the paths are required. export interface DiffusionTrainingStartRequest { base_model: string; + // Explicit family (sdxl / flux.1 / qwen-image / z-image). Optional: the backend + // resolves it from base_model when omitted, but the Train tab always sends it so a + // custom base still trains under the intended family. + model_family?: string | null; data_dir: string; output_dir: string; instance_prompt?: string | null; @@ -277,10 +281,18 @@ export interface DiffusionTrainingStartRequest { mixed_precision?: "bf16" | "fp16" | "no"; gradient_checkpointing?: boolean; lr_scheduler?: string; - // Forwarded to StableDiffusionXLPipeline.from_pretrained for a gated/private base repo. + // Forwarded to the pipeline's from_pretrained for a gated/private base repo (e.g. FLUX). hf_token?: string | null; } +// Paired step-indexed history arrays for the live loss + LR charts. `lr` entries may be +// null so a sparse learning-rate series still aligns with `steps` by index. +export interface DiffusionMetricHistory { + steps: number[]; + loss: number[]; + lr: Array; +} + // A snapshot of the current diffusion training job (GET /api/train/diffusion/status). export interface DiffusionTrainingStatus { active: boolean; @@ -298,6 +310,16 @@ export interface DiffusionTrainingStatus { lora_path: string | null; started_at: number | null; updated_at: number | null; + // Where the trained adapter was mirrored into the Studio LoRA catalog, and the family / + // base it was trained from -- lets the Train tab deploy the adapter onto the right base. + catalog_path?: string | null; + family?: string | null; + base_model?: string | null; + // Live throughput + peak VRAM (from the trainer's progress events). + samples_per_second?: number | null; + peak_memory_gb?: number | null; + // Bounded step/loss/lr history for the live charts. + metric_history?: DiffusionMetricHistory | null; } export async function startDiffusionTraining( @@ -328,11 +350,33 @@ export interface DiffusionDatasetSummary { caption_count: number; } +// Per-family training defaults (from GET /api/train/diffusion/info families[], added by +// the DiT-trainer backend). Absent on older backends; the Train tab falls back to a +// hardcoded family list when it is. +export interface DiffusionTrainableFamily { + name: string; + label: string; + default_base: string; + base_repos: string[]; + defaults?: { + lora_rank?: number; + learning_rate?: number; + resolution?: number; + train_steps?: number; + train_batch_size?: number; + mixed_precision?: "bf16" | "fp16" | "no"; + } | null; + vram_note?: string | null; + gated?: boolean | null; +} + // Where diffusion training reads/writes on this Studio, plus usable dataset folders. export interface DiffusionTrainingInfo { datasets_root: string; outputs_root: string; datasets: DiffusionDatasetSummary[]; + // Added by the multi-family trainer backend; tolerate its absence. + families?: DiffusionTrainableFamily[]; } export async function getDiffusionTrainingInfo(): Promise { @@ -357,3 +401,115 @@ export async function uploadDiffusionDataset( await authFetch("/api/train/diffusion/dataset", { method: "POST", body: form }), ); } + +// ── Dataset labeling + example imports (GET/PUT/DELETE .../dataset/{name}/...) ── +// One image in a training dataset folder, with its resolved caption. `caption_source` +// records where the caption came from ("metadata" beats a per-image "sidecar"; "none" +// when uncaptioned) so the labeling grid can highlight images that still need one. +export interface DiffusionDatasetImageRecord { + filename: string; + caption: string | null; + caption_source: "sidecar" | "metadata" | "none"; + width: number; + height: number; + size_bytes: number; +} + +export interface DiffusionDatasetImages { + name: string; + path: string; + images: DiffusionDatasetImageRecord[]; +} + +/** List every image in a dataset folder (including uncaptioned ones) for the grid. */ +export async function listDiffusionDatasetImages( + name: string, +): Promise { + return parseJson( + await authFetch(`/api/train/diffusion/dataset/${encodeURIComponent(name)}/images`), + ); +} + +/** Build the auth-protected thumbnail URL for a dataset image. Fetch it via + * fetchGalleryObjectUrl (Bearer auth) into an object URL; it can't be a plain . */ +export function diffusionDatasetImageUrl( + name: string, + filename: string, + thumb = 256, +): string { + const q = thumb > 0 ? `?thumb=${thumb}` : ""; + return `/api/train/diffusion/dataset/${encodeURIComponent(name)}/image/${encodeURIComponent(filename)}${q}`; +} + +/** Write (or, when blank, clear) a per-image caption sidecar. Returns the updated record. */ +export async function setDiffusionDatasetCaption( + name: string, + filename: string, + caption: string, +): Promise { + return parseJson( + await authFetch( + `/api/train/diffusion/dataset/${encodeURIComponent(name)}/caption/${encodeURIComponent(filename)}`, + { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ caption }), + }, + ), + ); +} + +/** Delete an image (and its caption + thumbnail) from a dataset folder. */ +export async function deleteDiffusionDatasetImage( + name: string, + filename: string, +): Promise { + const res = await authFetch( + `/api/train/diffusion/dataset/${encodeURIComponent(name)}/image/${encodeURIComponent(filename)}`, + { method: "DELETE" }, + ); + if (!res.ok) throw new Error(await readFastApiError(res)); +} + +// A curated, one-click-importable example image dataset. `license` is shown verbatim so +// users see the terms before importing; `suggested_trigger` seeds the trigger prompt. +export interface DiffusionDatasetExample { + id: string; + label: string; + repo: string; + description: string; + license: string; + image_cap: number; + suggested_trigger?: string | null; +} + +export async function listDiffusionDatasetExamples(): Promise { + const data = await parseJson<{ examples: DiffusionDatasetExample[] }>( + await authFetch("/api/train/diffusion/dataset-examples"), + ); + return data.examples; +} + +export interface DiffusionDatasetImportResult { + name: string; + path: string; + image_count: number; + caption_count: number; + imported: number; + license: string; + source_repo: string; +} + +/** Materialize a curated example dataset (by id) into a Studio dataset folder. */ +export async function importDiffusionDatasetExample( + id: string, + name?: string, +): Promise { + return parseJson( + await authFetch("/api/train/diffusion/dataset/import-example", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ id, name }), + }), + ); +} diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 54775145ce..e712f24c21 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -3,7 +3,6 @@ import { useCallback, useEffect, useMemo, useRef, useState, type ReactNode } from "react"; import { - AiMagicIcon, ArrowLeftRightIcon, ArrowReloadHorizontalIcon, Delete02Icon, @@ -32,6 +31,7 @@ import { import { Slider } from "@/components/ui/slider"; import { Spinner } from "@/components/ui/spinner"; import { Switch } from "@/components/ui/switch"; +import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { Textarea } from "@/components/ui/textarea"; import { InfoHint } from "@/components/ui/info-hint"; import { ModelSelector } from "@/components/assistant-ui/model-selector"; @@ -67,7 +67,7 @@ import { loadDiffusionModel, unloadDiffusionModel, } from "./api"; -import { DiffusionTrainDialog } from "./diffusion-train-dialog"; +import { DiffusionTrainPanel } from "./train/diffusion-train-panel"; // 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; @@ -464,34 +464,41 @@ function Field({ function AdvancedSelect({ label, hint, + desc, value, onValueChange, options, }: { label: string; hint?: ReactNode; + // A short always-visible description under the row (the hint tooltip carries the full + // detail). Used for controls whose label alone does not convey what they do. + desc?: string; value: string; onValueChange: (v: string) => void; options: Array<[string, string]>; }) { return ( -
- - {label} - {hint && {hint}} - - +
+
+ + {label} + {hint && {hint}} + + +
+ {desc &&

{desc}

}
); } @@ -948,8 +955,9 @@ 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); + // Page mode: "create" is the generation workspace; "train" is the full-page LoRA + // training workspace. Independent of the loaded generation model. + const [pageMode, setPageMode] = useState<"create" | "train">("create"); // 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); @@ -1025,6 +1033,9 @@ export function ImagesPage({ active = true }: { active?: boolean }) { // loaded, so the poll must roll the label back rather than advertise the failed // quant. `{ prev }` distinguishes "revert to null" from "nothing pending". const quantRevert = useRef<{ prev: string | null } | null>(null); + // A trained adapter awaiting deployment: after Deploy loads the base, the LoRA discovery + // effect applies this once the model is loaded + LoRA-capable for the matching family. + const pendingDeploy = useRef<{ loraId: string; family: string } | null>(null); const dismissLoadToast = useCallback(() => { if (loadToastId.current != null) toast.dismiss(loadToastId.current); @@ -1064,6 +1075,21 @@ export function ImagesPage({ active = true }: { active?: boolean }) { setLoras([]); } prevLoraFamilyRef.current = fam; + // A just-deployed adapter: now that the base is loaded + LoRA-capable, apply it (after + // the family-swap clear above so it isn't wiped). Only when the family matches what it + // was trained for; otherwise warn instead of silently applying an incompatible adapter. + const deploy = pendingDeploy.current; + if (deploy) { + pendingDeploy.current = null; + if (!deploy.family || deploy.family === fam) { + setLoras([{ id: deploy.loraId, weight: 1 }]); + } else { + toast.error( + `The trained adapter is for ${deploy.family}, but the loaded model is ` + + `${fam ?? "a different family"}, so it was not applied.`, + ); + } + } let cancelled = false; listDiffusionLoras(status?.family ?? undefined) .then((list) => { @@ -1565,6 +1591,36 @@ export function ImagesPage({ active = true }: { active?: boolean }) { [busy, handleLoad, quant], ); + // Deploy a freshly-trained adapter from the Train tab: switch to Create, load the base as + // a pipeline, and queue the adapter so the LoRA discovery effect applies it once the base + // is loaded + LoRA-capable. Seeds the prompt with the trigger phrase when provided. + const handleDeployAdapter = useCallback( + (args: { baseRepo: string; family: string; catalogPath: string; trigger: string }) => { + if (busy !== null) { + toast.error("Finish the current model load before deploying the adapter."); + return; + } + // The picker keys a local adapter by its filename stem (see diffusion_lora scan). + const base = args.catalogPath.replace(/\\/g, "/").split("/").pop() ?? ""; + const stem = base.replace(/\.(safetensors|gguf)$/i, ""); + if (!stem) { + toast.error("Could not resolve the trained adapter's name."); + return; + } + pendingDeploy.current = { loraId: stem, family: args.family }; + if (args.trigger.trim()) setPrompt(args.trigger.trim()); + setPageMode("create"); + setQuant(null); + const d = defaultsFor(args.baseRepo); + setSteps(d.steps); + setGuidance(d.guidance); + void handleLoad(args.baseRepo, { kind: "pipeline" }).then((started) => { + if (!started) pendingDeploy.current = null; + }); + }, + [busy, handleLoad], + ); + const handleUnload = useCallback(async () => { // Ejecting cancels any in-flight replacement load on the backend, so tear // down its client-side tracking too: the load poll reschedules on phase @@ -1815,7 +1871,8 @@ export function ImagesPage({ active = true }: { active?: boolean }) { to GGUF (or nothing loaded) and otherwise show why it is unavailable. */} {!status?.loaded || status.model_kind === "gguf" ? ( setTransformerQuant(v as typeof transformerQuant)} @@ -1829,9 +1886,11 @@ export function ImagesPage({ active = true }: { active?: boolean }) { ]} /> ) : ( -
- GGUF speed mode - GGUF models only +
+ + GGUF compute + + GGUF models only
)} setSelectorOpen(active && o)} />
- {/* Train a LoRA (SDXL): opens a self-contained dialog; available regardless of - whether a generation model is loaded. */} - + {/* Create | Train page-mode switch, next to the model selector. Create is the + generation workspace; Train is the full-page LoRA training workspace. */} + setPageMode(v as "create" | "train")}> + + + Create + + + Train + + + {/* 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. */} - + toggle, same icon in both states so it never moves). Highlighted when open. + Only meaningful in Create mode (load-time tuning), so hidden while training. */} + {pageMode === "create" && ( + + )}
- 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. ── */} + {/* Train mode: the full-page training workspace. Kept unmounted in Create mode so its + polling stops; Create's own state (gallery, model, workflow) is untouched. */} + {pageMode === "train" ? ( + setLoraRefreshKey((k) => k + 1)} + onDeploy={handleDeployAdapter} + /> + ) : ( + /* ── Controls rail + preview canvas. Padding mirrors the other tabs + (Export, Data Recipes): px-5 / sm:px-9, with a roomy bottom. ── */
{/* The controls rail. Plain card (the gray surface) with no header — the prompt + Generate button make the panel self-explanatory. */} @@ -2640,6 +2702,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
)}
+ )} ); } diff --git a/studio/frontend/src/features/images/train/dataset-labeling-grid.tsx b/studio/frontend/src/features/images/train/dataset-labeling-grid.tsx new file mode 100644 index 0000000000..48aa027c19 --- /dev/null +++ b/studio/frontend/src/features/images/train/dataset-labeling-grid.tsx @@ -0,0 +1,299 @@ +// 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 { ArrowLeft01Icon, ArrowRight01Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; + +import { Button } from "@/components/ui/button"; +import { Spinner } from "@/components/ui/spinner"; +import { Textarea } from "@/components/ui/textarea"; +import { cn } from "@/lib/utils"; +import { toast } from "@/lib/toast"; + +// One batch of images shown at a time in the labeling grid; larger sets page with < >. +const PAGE_SIZE = 24; + +import { + type DiffusionDatasetImageRecord, + deleteDiffusionDatasetImage, + diffusionDatasetImageUrl, + fetchGalleryObjectUrl, + listDiffusionDatasetImages, + setDiffusionDatasetCaption, +} from "../api"; + +// One tile: an auth-fetched thumbnail (object URL, revoked on unmount) plus a caption +// Textarea saved on blur. Uncaptioned tiles get a highlighted ring so a user labeling a +// small set can see at a glance what still needs a caption. +function LabelTile({ + dataset, + record, + onSaved, + onDeleted, +}: { + dataset: string; + record: DiffusionDatasetImageRecord; + onSaved: (rec: DiffusionDatasetImageRecord) => void; + onDeleted: (filename: string) => void; +}) { + const [thumb, setThumb] = useState(null); + const [caption, setCaption] = useState(record.caption ?? ""); + const [saving, setSaving] = useState(false); + const [savedTick, setSavedTick] = useState(false); + const [deleting, setDeleting] = useState(false); + // The last caption we persisted, so blur only writes when the text actually changed. + const persisted = useRef(record.caption ?? ""); + + // Load the thumbnail once; revoke the object URL on unmount to avoid a leak. + useEffect(() => { + let url: string | null = null; + let cancelled = false; + fetchGalleryObjectUrl(diffusionDatasetImageUrl(dataset, record.filename, 256)) + .then((u) => { + if (cancelled) { + URL.revokeObjectURL(u); + return; + } + url = u; + setThumb(u); + }) + .catch(() => { + /* a missing thumbnail just leaves the placeholder */ + }); + return () => { + cancelled = true; + if (url) URL.revokeObjectURL(url); + }; + }, [dataset, record.filename]); + + const save = useCallback(async () => { + const next = caption.trim(); + if (next === persisted.current.trim()) return; + setSaving(true); + try { + const updated = await setDiffusionDatasetCaption(dataset, record.filename, next); + persisted.current = updated.caption ?? ""; + setCaption(updated.caption ?? ""); + onSaved(updated); + setSavedTick(true); + window.setTimeout(() => setSavedTick(false), 1500); + } catch (e) { + toast.error(e instanceof Error ? e.message : "Failed to save caption"); + } finally { + setSaving(false); + } + }, [caption, dataset, record.filename, onSaved]); + + const remove = useCallback(async () => { + setDeleting(true); + try { + await deleteDiffusionDatasetImage(dataset, record.filename); + onDeleted(record.filename); + } catch (e) { + toast.error(e instanceof Error ? e.message : "Failed to delete image"); + setDeleting(false); + } + }, [dataset, record.filename, onDeleted]); + + const uncaptioned = caption.trim().length === 0; + + return ( +
+
+ {thumb ? ( + {record.filename} + ) : ( +
+ +
+ )} + +
+