From 0ebcdcefaf849d608b2ae3c15c21c5e2d41bd7c7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 15:50:37 +0000 Subject: [PATCH] Wire Create/Train tab switch into the Images page and deploy flow Replaces the Train LoRA dialog with a top-bar Create | Train segmented control next to the model selector. Create renders the existing generation workspace unchanged; Train renders the full-page training panel (unmounted in Create so its polling stops while the backend run and its retained metric history survive a tab switch). Adds a deploy handler: loading the trained adapter's base as a pipeline, queueing the adapter so the LoRA discovery effect applies it once the base is loaded and LoRA-capable for the matching family (with a mismatch warning), seeding the prompt with the trigger, and switching back to Create. Removes the now-unused dialog. --- .../images/diffusion-train-dialog.tsx | 510 ------------------ .../src/features/images/images-page.tsx | 148 +++-- 2 files changed, 101 insertions(+), 557 deletions(-) delete mode 100644 studio/frontend/src/features/images/diffusion-train-dialog.tsx diff --git a/studio/frontend/src/features/images/diffusion-train-dialog.tsx b/studio/frontend/src/features/images/diffusion-train-dialog.tsx deleted file mode 100644 index 25b260a03e..0000000000 --- a/studio/frontend/src/features/images/diffusion-train-dialog.tsx +++ /dev/null @@ -1,510 +0,0 @@ -// 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 54775145ce..6d22473377 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -32,6 +32,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 +68,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; @@ -948,8 +949,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 +1027,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 +1069,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 +1585,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 @@ -1910,54 +1960,57 @@ export function ImagesPage({ active = true }: { active?: boolean }) { onOpenChange={(o) => 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 +2693,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
)} + )} ); }