ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Offer example datasets in the Train dropdown with previews

Browse source 
Add an Examples group to the training-images dropdown that imports a
curated dataset in one pick, alongside the existing cards. Cards now show
up to three preview thumbnails pulled from the public HF datasets-server
so the set is visible before download. Hide the trigger prompt when every
image already has a caption (a captioned style set needs no trigger), and
turn the training-settings toggle into a ghost button with a rotating
chevron.
This commit is contained in:
Daniel Han 2026-07-03 06:27:19 +00:00
commit 0a3ccd844e
2 changed files with 232 additions and 96 deletions
Download patch file Download diff file Expand all files Collapse all files

147
studio/frontend/src/features/images/train/diffusion-train-panel.tsx
View file

@ -3,6 +3,9 @@
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import { ArrowDown01Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Label } from "@/components/ui/label";
@ -13,11 +16,13 @@ import { cn } from "@/lib/utils";
import { toast } from "@/lib/toast";
import {
type DiffusionDatasetExample,
type DiffusionTrainableFamily,
type DiffusionTrainingInfo,
type DiffusionTrainingStatus,
getDiffusionTrainingInfo,
getDiffusionTrainingStatus,
listDiffusionDatasetExamples,
startDiffusionTraining,
stopDiffusionTraining,
uploadDiffusionDataset,
@ -25,7 +30,7 @@ import {
import { DatasetLabelingGrid, LabelingGridToggle } from "./dataset-labeling-grid";
import { DatasetShowcase } from "./dataset-showcase";
import { DiffusionCharts } from "./diffusion-charts";
import { ExampleDatasetCards } from "./example-dataset-cards";
import { ExampleDatasetCards, runExampleImport } from "./example-dataset-cards";
// The families the Train tab can train, in the popularity order the user asked for. This is
// the fallback used when the backend's /info does not yet report families (older backend);
@ -73,6 +78,8 @@ const FAMILY_PRESETS: FamilyPreset[] = [
const CUSTOM_BASE = "__custom__";
const UPLOAD_DATASET = "__upload__";
// Dataset-select option value prefix for a not-yet-imported example; picking it imports.
const EXAMPLE_PREFIX = "example:";
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";
@ -159,6 +166,8 @@ export function DiffusionTrainPanel({
const fileInputRef = useRef<HTMLInputElement | null>(null);
const [gridOpen, setGridOpen] = useState(false);
const [gridRefresh, setGridRefresh] = useState(0);
const [examples, setExamples] = useState<DiffusionDatasetExample[]>([]);
const [importingId, setImportingId] = useState<string | null>(null);
const [outputDir, setOutputDir] = useState("");
const [instancePrompt, setInstancePrompt] = useState("");
@ -199,6 +208,58 @@ export function DiffusionTrainPanel({
});
}, [active, refreshInfo]);
// Load the curated example list once (for the dropdown group + the cards). Best-effort:
// an older backend without the endpoint just yields no examples.
useEffect(() => {
if (!active) return;
let cancelled = false;
listDiffusionDatasetExamples()
.then((list) => {
if (!cancelled) setExamples(list);
})
.catch(() => {
if (!cancelled) setExamples([]);
});
return () => {
cancelled = true;
};
}, [active]);
// Examples whose folder is not on disk yet: shown in the dropdown's Examples group and as
// cards. An example imports into a folder named after its id, so a matching dataset name
// means it is already imported (and appears as a normal dataset instead).
const importedNames = useMemo(
() => new Set((info?.datasets ?? []).map((d) => d.name)),
[info?.datasets],
);
const pendingExamples = useMemo(
() => examples.filter((ex) => !importedNames.has(ex.id)),
[examples, importedNames],
);
// Import a curated example, then select the resulting folder. Seeds the trigger prompt from
// the example only when the field is meaningful (the import has no captions of its own).
const importExample = useCallback(
async (ex: DiffusionDatasetExample) => {
setImportingId(ex.id);
try {
const res = await runExampleImport(ex);
await refreshInfo();
setDataset(res.name);
setGridOpen(false);
setGridRefresh((k) => k + 1);
if (ex.suggested_trigger && res.caption_count === 0 && !instancePrompt.trim()) {
setInstancePrompt(ex.suggested_trigger);
}
} catch (e) {
toast.error(e instanceof Error ? e.message : "Import failed");
} finally {
setImportingId(null);
}
},
[refreshInfo, instancePrompt],
);
// If the loaded generation model is a trainable family, jump the family selector to it
// once (only when the panel first sees a loaded family).
const seededFromLoaded = useRef(false);
@ -267,6 +328,13 @@ export function DiffusionTrainPanel({
const selectedDataset =
dataset !== UPLOAD_DATASET ? info?.datasets.find((d) => d.name === dataset) : undefined;
// A dataset where every image already ships a caption needs no trigger prompt; hide the
// field and explain why. Partial/no captions (or upload mode) still show it.
const fullyCaptioned = Boolean(
selectedDataset &&
selectedDataset.image_count > 0 &&
selectedDataset.caption_count >= selectedDataset.image_count,
);
// Map the backend's paired history arrays into the chart component's {step,value} series.
const lossHistory: TrainingSeriesPoint[] = useMemo(() => {
@ -494,11 +562,18 @@ export function DiffusionTrainPanel({
<select
value={dataset}
onChange={(e) => {
setDataset(e.target.value);
const v = e.target.value;
if (v.startsWith(EXAMPLE_PREFIX)) {
const ex = pendingExamples.find((x) => x.id === v.slice(EXAMPLE_PREFIX.length));
if (ex) void importExample(ex);
return; // controlled select snaps back to the current dataset while importing
}
setDataset(v);
setGridOpen(false);
}}
className={selectClass}
aria-label="Training images"
disabled={importingId !== null}
>
{(info?.datasets ?? []).map((d) => (
<option key={d.name} value={d.name}>
@ -506,8 +581,22 @@ export function DiffusionTrainPanel({
{d.caption_count > 0 ? `, ${d.caption_count} captions` : ""})
</option>
))}
{pendingExamples.length > 0 && (
<optgroup label="Examples (one-click import)">
{pendingExamples.map((ex) => (
<option key={ex.id} value={`${EXAMPLE_PREFIX}${ex.id}`}>
{ex.label} ({ex.image_cap} images, {ex.license})
</option>
))}
</optgroup>
)}
<option value={UPLOAD_DATASET}>Upload new images...</option>
</select>
{importingId && (
<p className="text-[11px] text-muted-foreground">
Importing {examples.find((e) => e.id === importingId)?.label ?? "example"}...
</p>
)}
{dataset === UPLOAD_DATASET ? (
<div className="grid gap-1.5 rounded-md border border-dashed border-border p-2">
@ -578,14 +667,9 @@ export function DiffusionTrainPanel({
)}
<ExampleDatasetCards
onImported={(res, ex) => {
void refreshInfo();
setDataset(res.name);
setGridRefresh((k) => k + 1);
if (ex.suggested_trigger && !instancePrompt.trim()) {
setInstancePrompt(ex.suggested_trigger);
}
}}
examples={pendingExamples}
busyId={importingId}
onImport={(ex) => void importExample(ex)}
/>
</div>
@ -600,27 +684,40 @@ export function DiffusionTrainPanel({
className="h-8 text-xs"
/>
</div>
<div className="grid gap-1.5">
<Label className="text-xs">Trigger prompt (how you&apos;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>
{fullyCaptioned ? (
<p className="text-[11px] leading-snug text-muted-foreground">
All {selectedDataset?.image_count} images have captions - no trigger prompt needed.
The style applies to any prompt after training.
</p>
) : (
<div className="grid gap-1.5">
<Label className="text-xs">
Trigger prompt (how you&apos;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>
)}
{/* Collapsed training settings */}
<button
<Button
type="button"
className="w-fit text-xs text-muted-foreground underline-offset-2 hover:underline"
variant="ghost"
size="sm"
className="h-7 w-fit gap-1.5 px-1 text-xs text-muted-foreground hover:text-foreground"
onClick={() => setShowAdvanced((s) => !s)}
aria-expanded={showAdvanced}
>
{showAdvanced
? "Hide training settings"
: "Training settings (defaults suit a first run)"}
</button>
<HugeiconsIcon
icon={ArrowDown01Icon}
className={cn("size-3.5 transition-transform", showAdvanced && "rotate-180")}
/>
{showAdvanced ? "Training settings" : "Training settings (defaults suit a first run)"}
</Button>
{showAdvanced && (
<>
<div className="grid grid-cols-2 gap-3 sm:grid-cols-4">

181
studio/frontend/src/features/images/train/example-dataset-cards.tsx
View file

@ -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 { useEffect, useState } from "react";
import { Button } from "@/components/ui/button";
import { toast } from "@/lib/toast";
@ -10,62 +10,87 @@ import {
type DiffusionDatasetExample,
type DiffusionDatasetImportResult,
importDiffusionDatasetExample,
listDiffusionDatasetExamples,
} from "../api";
// One-click example-dataset importers. Each card shows the license so users see the terms
// before importing. On success the parent refreshes its dataset list and selects the
// imported folder (and can seed the trigger prompt from suggested_trigger).
//
// Layout: one card per row (the config column is only ~340px, so a two-column grid wrapped
// titles one word per line and let the long license text overrun into the next card). The
// license is a compact truncated badge with the full text in its title tooltip.
export function ExampleDatasetCards({
onImported,
}: {
onImported: (
result: DiffusionDatasetImportResult,
example: DiffusionDatasetExample,
) => void;
}) {
const [examples, setExamples] = useState<DiffusionDatasetExample[] | null>(null);
const [busyId, setBusyId] = useState<string | null>(null);
// Best-effort preview thumbnails pulled from the public HF datasets-server. Cached per repo
// (module-level) so re-renders and re-mounts do not refetch. A repo that the server cannot
// serve (e.g. diffusers/dog-example) resolves to an empty list and the card renders without
// previews - the import still works.
const _previewCache = new Map<string, Promise<string[]>>();
async function fetchPreviews(repo: string): Promise<string[]> {
const cached = _previewCache.get(repo);
if (cached) return cached;
const p = (async () => {
try {
const res = await fetch(
`https://datasets-server.huggingface.co/first-rows?dataset=${encodeURIComponent(
repo,
)}&config=default&split=train`,
);
if (!res.ok) return [];
const data = (await res.json()) as {
features?: { name: string; type?: { _type?: string } }[];
rows?: { row: Record<string, unknown> }[];
};
const imageCol = data.features?.find((f) => f.type?._type === "Image")?.name;
if (!imageCol || !data.rows) return [];
const urls: string[] = [];
for (const r of data.rows) {
const cell = r.row[imageCol] as { src?: string } | undefined;
if (cell?.src) urls.push(cell.src);
if (urls.length >= 3) break;
}
return urls;
} catch {
return [];
}
})();
_previewCache.set(repo, p);
return p;
}
function ExamplePreviews({ repo }: { repo: string }) {
const [urls, setUrls] = useState<string[] | null>(null);
useEffect(() => {
let cancelled = false;
listDiffusionDatasetExamples()
.then((list) => {
if (!cancelled) setExamples(list);
})
.catch(() => {
if (!cancelled) setExamples([]); // older backend: just hide the cards
});
void fetchPreviews(repo).then((u) => {
if (!cancelled) setUrls(u);
});
return () => {
cancelled = true;
};
}, []);
}, [repo]);
const doImport = useCallback(
async (ex: DiffusionDatasetExample) => {
setBusyId(ex.id);
try {
const res = await importDiffusionDatasetExample(ex.id);
toast.success(
res.imported > 0
? `Imported ${res.image_count} images into "${res.name}"`
: `"${res.name}" already imported (${res.image_count} images)`,
);
onImported(res, ex);
} catch (e) {
toast.error(e instanceof Error ? e.message : "Import failed");
} finally {
setBusyId(null);
}
},
[onImported],
if (!urls || urls.length === 0) return null;
return (
<div className="flex gap-1">
{urls.map((u) => (
<div key={u} className="size-10 shrink-0 overflow-hidden rounded-md bg-muted">
<img src={u} alt="" loading="lazy" className="size-full object-cover" />
</div>
))}
</div>
);
}
if (!examples || examples.length === 0) return null;
// One-click example-dataset importers. Each card shows the license so users see the terms
// before importing, plus a few preview thumbnails so the set is visible before download. On
// success the parent refreshes its dataset list and selects the imported folder (and can
// seed the trigger prompt from suggested_trigger).
//
// Layout: one card per row (the config column is narrow, so a two-column grid wrapped titles
// one word per line and let the long license text overrun into the next card).
export function ExampleDatasetCards({
examples,
busyId,
onImport,
}: {
examples: DiffusionDatasetExample[];
busyId: string | null;
onImport: (ex: DiffusionDatasetExample) => void;
}) {
if (examples.length === 0) return null;
return (
<div className="grid gap-2">
@ -76,37 +101,51 @@ export function ExampleDatasetCards({
{examples.map((ex) => (
<div
key={ex.id}
className="flex min-w-0 items-center gap-3 rounded-lg border border-border p-2.5"
className="flex min-w-0 flex-col gap-2 rounded-lg border border-border p-2.5"
>
<div className="flex min-w-0 flex-1 flex-col gap-1">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 flex-1 truncate text-xs font-medium">
{ex.label}
</span>
<span
className="max-w-[110px] shrink truncate rounded-full bg-secondary px-2 py-0.5 text-[10px] font-normal text-secondary-foreground"
title={ex.license}
>
{ex.license}
</span>
<div className="flex min-w-0 items-center gap-3">
<div className="flex min-w-0 flex-1 flex-col gap-1">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 flex-1 truncate text-xs font-medium">{ex.label}</span>
<span
className="max-w-[110px] shrink truncate rounded-full bg-secondary px-2 py-0.5 text-[10px] font-normal text-secondary-foreground"
title={ex.license}
>
{ex.license}
</span>
</div>
<p className="line-clamp-2 text-[11px] leading-snug text-muted-foreground">
{ex.description}
</p>
</div>
<p className="line-clamp-2 text-[11px] leading-snug text-muted-foreground">
{ex.description}
</p>
<Button
type="button"
size="sm"
variant="secondary"
className="h-7 shrink-0 self-center px-3 text-xs"
onClick={() => onImport(ex)}
disabled={busyId !== null}
>
{busyId === ex.id ? "Importing..." : "Import"}
</Button>
</div>
<Button
type="button"
size="sm"
variant="secondary"
className="h-7 shrink-0 self-center px-3 text-xs"
onClick={() => void doImport(ex)}
disabled={busyId !== null}
>
{busyId === ex.id ? "Importing..." : "Import"}
</Button>
<ExamplePreviews repo={ex.repo} />
</div>
))}
</div>
</div>
);
}
// Shared import helper so the panel's dropdown and the cards import identically.
export async function runExampleImport(
ex: DiffusionDatasetExample,
): Promise<DiffusionDatasetImportResult> {
const res = await importDiffusionDatasetExample(ex.id);
toast.success(
res.imported > 0
? `Imported ${res.image_count} images into "${res.name}"`
: `"${res.name}" already imported (${res.image_count} images)`,
);
return res;
}