diff --git a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx index 5706aa2ad6..8c4ce6b7c0 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -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(null); const [gridOpen, setGridOpen] = useState(false); const [gridRefresh, setGridRefresh] = useState(0); + const [examples, setExamples] = useState([]); + const [importingId, setImportingId] = useState(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({ + {importingId && ( +

+ Importing {examples.find((e) => e.id === importingId)?.label ?? "example"}... +

+ )} {dataset === UPLOAD_DATASET ? (
@@ -578,14 +667,9 @@ export function DiffusionTrainPanel({ )} { - 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)} />
@@ -600,27 +684,40 @@ export function DiffusionTrainPanel({ className="h-8 text-xs" /> -
- - setInstancePrompt(e.target.value)} - className="h-8 text-xs" - /> -
+ {fullyCaptioned ? ( +

+ All {selectedDataset?.image_count} images have captions - no trigger prompt needed. + The style applies to any prompt after training. +

+ ) : ( +
+ + setInstancePrompt(e.target.value)} + className="h-8 text-xs" + /> +
+ )} {/* Collapsed training settings */} - + + {showAdvanced ? "Training settings" : "Training settings (defaults suit a first run)"} + {showAdvanced && ( <>
diff --git a/studio/frontend/src/features/images/train/example-dataset-cards.tsx b/studio/frontend/src/features/images/train/example-dataset-cards.tsx index 4a80ca53a2..8bf5d3d38c 100644 --- a/studio/frontend/src/features/images/train/example-dataset-cards.tsx +++ b/studio/frontend/src/features/images/train/example-dataset-cards.tsx @@ -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(null); - const [busyId, setBusyId] = useState(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>(); +async function fetchPreviews(repo: string): Promise { + 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 }[]; + }; + 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(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 ( +
+ {urls.map((u) => ( +
+ +
+ ))} +
); +} - 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 (
@@ -76,37 +101,51 @@ export function ExampleDatasetCards({ {examples.map((ex) => (
-
-
- - {ex.label} - - - {ex.license} - +
+
+
+ {ex.label} + + {ex.license} + +
+

+ {ex.description} +

-

- {ex.description} -

+
- +
))}
); } + +// Shared import helper so the panel's dropdown and the cards import identically. +export async function runExampleImport( + ex: DiffusionDatasetExample, +): Promise { + 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; +}