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.
This commit is contained in:
Daniel Han
parent
7cd8f28310
commit
0ebcdcefaf
2 changed files with 101 additions and 557 deletions
|
|
@ -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<DiffusionTrainingInfo | null>(null);
|
||||
const [dataset, setDataset] = useState<string>(UPLOAD_DATASET);
|
||||
const [uploadName, setUploadName] = useState("my-images");
|
||||
const [uploading, setUploading] = useState(false);
|
||||
const fileInputRef = useRef<HTMLInputElement | null>(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<DiffusionTrainingStatus | null>(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<DiffusionTrainingInfo | null> => {
|
||||
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 (
|
||||
<Dialog open={open} onOpenChange={onOpenChange}>
|
||||
<DialogContent className="flex max-h-[85vh] max-w-lg flex-col">
|
||||
<DialogHeader>
|
||||
<DialogTitle>Train an SDXL LoRA</DialogTitle>
|
||||
<DialogDescription>
|
||||
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.
|
||||
</DialogDescription>
|
||||
</DialogHeader>
|
||||
|
||||
<div className="grid gap-3 overflow-y-auto py-2 pr-1">
|
||||
{/* 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. */}
|
||||
<div className="grid gap-1.5">
|
||||
<Label className="text-xs">Base model to train on</Label>
|
||||
<select
|
||||
value={baseChoice}
|
||||
onChange={(e) => setBaseChoice(e.target.value)}
|
||||
className={selectClass}
|
||||
>
|
||||
{SDXL_BASES.map((b) => (
|
||||
<option key={b.id} value={b.id}>
|
||||
{b.label}
|
||||
</option>
|
||||
))}
|
||||
{defaultBaseModel && !SDXL_BASES.some((b) => b.id === defaultBaseModel) && (
|
||||
<option value={defaultBaseModel}>Loaded model: {defaultBaseModel}</option>
|
||||
)}
|
||||
<option value={CUSTOM_BASE}>Custom SDXL repo or local path...</option>
|
||||
</select>
|
||||
{baseChoice === CUSTOM_BASE && (
|
||||
<Input
|
||||
value={customBase}
|
||||
placeholder="my-org/my-sdxl-finetune or /path/to/sdxl-pipeline"
|
||||
spellCheck={false}
|
||||
onChange={(e) => setCustomBase(e.target.value)}
|
||||
className="h-8 text-xs"
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{/* 2. Training images: pick an existing dataset folder or upload straight from
|
||||
the browser - no shell access or Studio-home knowledge needed. */}
|
||||
<div className="grid gap-1.5">
|
||||
<Label className="text-xs">Training images</Label>
|
||||
<select
|
||||
value={dataset}
|
||||
onChange={(e) => setDataset(e.target.value)}
|
||||
className={selectClass}
|
||||
>
|
||||
{(info?.datasets ?? []).map((d) => (
|
||||
<option key={d.name} value={d.name}>
|
||||
{d.name} ({d.image_count} image{d.image_count === 1 ? "" : "s"}
|
||||
{d.caption_count > 0 ? `, ${d.caption_count} captions` : ""})
|
||||
</option>
|
||||
))}
|
||||
<option value={UPLOAD_DATASET}>Upload new images...</option>
|
||||
</select>
|
||||
{dataset === UPLOAD_DATASET && (
|
||||
<div className="grid gap-1.5 rounded-md border border-dashed border-border p-2">
|
||||
<Input
|
||||
value={uploadName}
|
||||
placeholder="my-style-photos"
|
||||
spellCheck={false}
|
||||
onChange={(e) => setUploadName(e.target.value)}
|
||||
className="h-8 text-xs"
|
||||
aria-label="New dataset name"
|
||||
/>
|
||||
<div className="flex items-center gap-2">
|
||||
<input
|
||||
ref={fileInputRef}
|
||||
type="file"
|
||||
multiple
|
||||
accept={DATASET_FILE_ACCEPT}
|
||||
className="min-w-0 flex-1 text-xs file:mr-2 file:rounded-md file:border-0 file:bg-muted file:px-2 file:py-1 file:text-xs"
|
||||
aria-label="Training images"
|
||||
/>
|
||||
<Button
|
||||
type="button"
|
||||
size="sm"
|
||||
variant="secondary"
|
||||
className="h-8 shrink-0"
|
||||
onClick={onUpload}
|
||||
disabled={uploading}
|
||||
>
|
||||
{uploading ? "Uploading..." : "Upload"}
|
||||
</Button>
|
||||
</div>
|
||||
<p className="text-[11px] text-muted-foreground">
|
||||
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.
|
||||
</p>
|
||||
</div>
|
||||
)}
|
||||
{selectedDataset && selectedDataset.caption_count === 0 && (
|
||||
<p className="text-[11px] text-muted-foreground">
|
||||
No caption files in this dataset - the trigger prompt below will be used
|
||||
as the caption for every image.
|
||||
</p>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{/* 3. What to call the result + how to trigger it. */}
|
||||
<div className="grid gap-1.5">
|
||||
<Label className="text-xs">Adapter name</Label>
|
||||
<Input
|
||||
value={outputDir}
|
||||
placeholder="my-style-lora"
|
||||
spellCheck={false}
|
||||
onChange={(e) => setOutputDir(e.target.value)}
|
||||
className="h-8 text-xs"
|
||||
/>
|
||||
</div>
|
||||
<div className="grid gap-1.5">
|
||||
<Label className="text-xs">Trigger prompt (how you'll invoke the style later)</Label>
|
||||
<Input
|
||||
value={instancePrompt}
|
||||
placeholder="a photo in SKS style"
|
||||
onChange={(e) => setInstancePrompt(e.target.value)}
|
||||
className="h-8 text-xs"
|
||||
/>
|
||||
</div>
|
||||
|
||||
{/* 4. Hyperparameters, collapsed: the defaults suit a first run, and hiding
|
||||
them keeps the primary flow at three decisions. */}
|
||||
<button
|
||||
type="button"
|
||||
className="w-fit text-xs text-muted-foreground underline-offset-2 hover:underline"
|
||||
onClick={() => setShowAdvanced((s) => !s)}
|
||||
aria-expanded={showAdvanced}
|
||||
>
|
||||
{showAdvanced ? "Hide training settings" : "Training settings (defaults suit a first run)"}
|
||||
</button>
|
||||
{showAdvanced && (
|
||||
<>
|
||||
<div className="grid grid-cols-2 gap-3 sm:grid-cols-4">
|
||||
<div className="grid gap-1.5">
|
||||
<Label className="text-xs">Steps</Label>
|
||||
<Input
|
||||
type="number"
|
||||
min={1}
|
||||
value={steps}
|
||||
onChange={(e) => setSteps(Number(e.target.value) || 1)}
|
||||
className="h-8 text-xs"
|
||||
/>
|
||||
</div>
|
||||
<div className="grid gap-1.5">
|
||||
<Label className="text-xs">LoRA rank</Label>
|
||||
<Input
|
||||
type="number"
|
||||
min={1}
|
||||
value={rank}
|
||||
onChange={(e) => setRank(Number(e.target.value) || 1)}
|
||||
className="h-8 text-xs"
|
||||
/>
|
||||
</div>
|
||||
<div className="grid gap-1.5">
|
||||
<Label className="text-xs">Resolution</Label>
|
||||
<Input
|
||||
type="number"
|
||||
min={64}
|
||||
step={64}
|
||||
value={resolution}
|
||||
onChange={(e) => setResolution(Number(e.target.value) || 1024)}
|
||||
className="h-8 text-xs"
|
||||
/>
|
||||
</div>
|
||||
<div className="grid gap-1.5">
|
||||
<Label className="text-xs">Batch</Label>
|
||||
<Input
|
||||
type="number"
|
||||
min={1}
|
||||
value={batchSize}
|
||||
onChange={(e) => setBatchSize(Number(e.target.value) || 1)}
|
||||
className="h-8 text-xs"
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
<div className="grid grid-cols-2 gap-3">
|
||||
<div className="grid gap-1.5">
|
||||
<Label className="text-xs">Learning rate</Label>
|
||||
<Input
|
||||
type="number"
|
||||
step={0.00001}
|
||||
min={0}
|
||||
value={learningRate}
|
||||
onChange={(e) => setLearningRate(Number(e.target.value) || 0.0001)}
|
||||
className="h-8 text-xs"
|
||||
/>
|
||||
</div>
|
||||
<div className="grid gap-1.5">
|
||||
<Label className="text-xs">Precision</Label>
|
||||
<select
|
||||
value={precision}
|
||||
onChange={(e) => setPrecision(e.target.value as "bf16" | "fp16" | "no")}
|
||||
className={selectClass}
|
||||
>
|
||||
<option value="bf16">bf16 (default)</option>
|
||||
<option value="fp16">fp16 (older GPUs)</option>
|
||||
<option value="no">fp32 (no mixed)</option>
|
||||
</select>
|
||||
</div>
|
||||
</div>
|
||||
</>
|
||||
)}
|
||||
|
||||
{status && status.status !== "idle" && (
|
||||
<div className="rounded-lg border border-border bg-muted/30 p-3 text-xs">
|
||||
<div className="mb-1 flex items-center justify-between">
|
||||
<span className="font-medium capitalize">{status.status}</span>
|
||||
<span className="text-muted-foreground">
|
||||
{status.total_steps > 0 ? `${status.step}/${status.total_steps}` : ""}
|
||||
</span>
|
||||
</div>
|
||||
<div className="mb-2 h-1.5 w-full overflow-hidden rounded-full bg-border">
|
||||
<div className="h-full bg-primary transition-all" style={{ width: `${pct}%` }} />
|
||||
</div>
|
||||
<div className="text-muted-foreground">
|
||||
{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 && (
|
||||
<div className="mt-1 break-all text-[11px]">Saved: {status.lora_path}</div>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<DialogFooter>
|
||||
{active ? (
|
||||
<Button type="button" variant="destructive" onClick={onStop}>
|
||||
Stop
|
||||
</Button>
|
||||
) : completed ? (
|
||||
<>
|
||||
<Button type="button" variant="secondary" onClick={() => onOpenChange(false)}>
|
||||
Done - open the LoRA picker
|
||||
</Button>
|
||||
<Button type="button" onClick={onStart} disabled={starting || uploading}>
|
||||
{starting ? "Starting..." : "Train another"}
|
||||
</Button>
|
||||
</>
|
||||
) : (
|
||||
<Button type="button" onClick={onStart} disabled={starting || uploading}>
|
||||
{starting ? "Starting..." : "Start training"}
|
||||
</Button>
|
||||
)}
|
||||
</DialogFooter>
|
||||
</DialogContent>
|
||||
</Dialog>
|
||||
);
|
||||
}
|
||||
|
|
@ -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<LoraSpecInput[]>([]);
|
||||
const [availableLoras, setAvailableLoras] = useState<DiffusionLoraInfo[]>([]);
|
||||
// "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)}
|
||||
/>
|
||||
<div className="flex items-center gap-2">
|
||||
{/* Train a LoRA (SDXL): opens a self-contained dialog; available regardless of
|
||||
whether a generation model is loaded. */}
|
||||
<Button
|
||||
type="button"
|
||||
variant="ghost"
|
||||
size="sm"
|
||||
className="h-[34px]"
|
||||
onClick={() => setTrainOpen(true)}
|
||||
title="Teach SDXL your own style or subject from a folder of images"
|
||||
>
|
||||
<HugeiconsIcon icon={AiMagicIcon} className="mr-1.5 size-3.5" />
|
||||
Train LoRA
|
||||
</Button>
|
||||
{/* Create | Train page-mode switch, next to the model selector. Create is the
|
||||
generation workspace; Train is the full-page LoRA training workspace. */}
|
||||
<Tabs value={pageMode} onValueChange={(v) => setPageMode(v as "create" | "train")}>
|
||||
<TabsList className="h-[34px]">
|
||||
<TabsTrigger value="create">Create</TabsTrigger>
|
||||
<TabsTrigger value="train">
|
||||
<HugeiconsIcon icon={AiMagicIcon} className="mr-1 size-3.5" />
|
||||
Train
|
||||
</TabsTrigger>
|
||||
</TabsList>
|
||||
</Tabs>
|
||||
{/* 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. */}
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => setAdvancedOpen((o) => !o)}
|
||||
aria-label={advancedOpen ? "Hide advanced options" : "Show advanced options"}
|
||||
aria-pressed={advancedOpen}
|
||||
title="Advanced options"
|
||||
className={cn(
|
||||
"flex h-[34px] w-[34px] items-center justify-center rounded-full transition-colors focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
|
||||
advancedOpen
|
||||
? "bg-muted text-foreground"
|
||||
: "text-muted-foreground hover:bg-muted hover:text-foreground",
|
||||
)}
|
||||
>
|
||||
<HugeiconsIcon icon={LayoutAlignRightIcon} className="size-4" />
|
||||
</button>
|
||||
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" && (
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => setAdvancedOpen((o) => !o)}
|
||||
aria-label={advancedOpen ? "Hide advanced options" : "Show advanced options"}
|
||||
aria-pressed={advancedOpen}
|
||||
title="Advanced options"
|
||||
className={cn(
|
||||
"flex h-[34px] w-[34px] items-center justify-center rounded-full transition-colors focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
|
||||
advancedOpen
|
||||
? "bg-muted text-foreground"
|
||||
: "text-muted-foreground hover:bg-muted hover:text-foreground",
|
||||
)}
|
||||
>
|
||||
<HugeiconsIcon icon={LayoutAlignRightIcon} className="size-4" />
|
||||
</button>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
<DiffusionTrainDialog
|
||||
open={active && trainOpen}
|
||||
onOpenChange={setTrainOpen}
|
||||
defaultBaseModel={
|
||||
status?.family === "sdxl"
|
||||
? // Prefer base_repo (the full diffusers pipeline) over repo_id: for a GGUF or
|
||||
// single-file SDXL load repo_id is the checkpoint path, which the trainer's
|
||||
// from_pretrained cannot open. base_repo is the companion pipeline.
|
||||
status?.base_repo ?? status?.repo_id ?? undefined
|
||||
: undefined
|
||||
}
|
||||
onTrainingComplete={() => 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" ? (
|
||||
<DiffusionTrainPanel
|
||||
active={active && pageMode === "train"}
|
||||
loadedFamily={status?.family ?? null}
|
||||
loadedBaseRepo={
|
||||
// Prefer base_repo (the full diffusers pipeline) over repo_id: for a GGUF or
|
||||
// single-file load repo_id is the checkpoint path, not a trainable base.
|
||||
status?.base_repo ?? status?.repo_id ?? null
|
||||
}
|
||||
onTrainingComplete={() => 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. ── */
|
||||
<div className="flex min-h-0 min-w-0 flex-1 gap-4 overflow-hidden px-5 pb-8 sm:px-9">
|
||||
{/* 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 }) {
|
|||
</div>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue