diff --git a/studio/frontend/src/features/images/train/diffusion-charts.tsx b/studio/frontend/src/features/images/train/diffusion-charts.tsx new file mode 100644 index 0000000000..c4378156cf --- /dev/null +++ b/studio/frontend/src/features/images/train/diffusion-charts.tsx @@ -0,0 +1,143 @@ +// 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 { type ReactElement, useMemo } from "react"; + +import type { TrainingSeriesPoint } from "@/features/training"; +// The loss + LR cards are pure presentational (props only), so reuse them directly. We do +// NOT reuse ChartsSection/ChartsContent: those also render Grad Norm and an Eval Loss card, +// which are meaningless for diffusion LoRA training and showed as an empty card and an +// "Evaluation not configured" placeholder. This is a diffusion-only two-card layout. +// eslint-disable-next-line no-restricted-imports +import { LearningRateChartCard } from "@/features/studio/sections/charts/learning-rate-chart-card"; +// eslint-disable-next-line no-restricted-imports +import { TrainingLossChartCard } from "@/features/studio/sections/charts/training-loss-chart-card"; +// eslint-disable-next-line no-restricted-imports +import { + MAX_RENDER_POINTS, + buildStepTicks, + buildYDomain, + compressSeries, + ema, +} from "@/features/studio/sections/charts/utils"; + +// Fixed presentation defaults (the LLM tab exposes these via a settings sheet; here we pick +// sensible constants): EMA smoothing on, linear scale, raw + smoothed + average lines shown, +// no outlier trimming (diffusion loss is naturally noisy, not spiky-with-outliers). +const SMOOTHING = 0.8; + +function toLossItems(series: TrainingSeriesPoint[]): { step: number; loss: number }[] { + return series + .filter((p) => Number.isFinite(p.value)) + .map((p) => ({ step: p.step, loss: p.value })); +} + +// The x-domain that spans all points (the LLM tab supports a scrollable window; a training +// run here is short enough to always show the whole thing). +function fullStepDomain(steps: number[]): [number, number] { + if (steps.length === 0) return [0, 1]; + const min = steps[0]; + const max = steps[steps.length - 1]; + if (min === max) return [min, min + 4]; + if (max - min < 6) return [Math.max(0, max - 6), max]; + return [min, max]; +} + +// A diffusion-only metrics view: just Training Loss and Learning Rate, side by side, with a +// note under the loss card explaining why per-step loss looks noisy. +export function DiffusionCharts({ + lossHistory, + lrHistory, +}: { + lossHistory: TrainingSeriesPoint[]; + lrHistory: TrainingSeriesPoint[]; +}): ReactElement | null { + const lossItems = useMemo(() => toLossItems(lossHistory), [lossHistory]); + const smoothed = useMemo( + () => (lossItems.length > 0 ? ema(lossItems, SMOOTHING) : []), + [lossItems], + ); + const reducedLoss = useMemo( + () => compressSeries(smoothed, MAX_RENDER_POINTS), + [smoothed], + ); + const lossData = useMemo( + () => + reducedLoss.map((p) => ({ + ...p, + displayLoss: p.loss, + displaySmoothed: p.smoothed, + })), + [reducedLoss], + ); + + const lrData = useMemo( + () => + compressSeries( + lrHistory + .filter((p) => Number.isFinite(p.value)) + .map((p) => ({ step: p.step, lr: p.value, displayLr: p.value })), + MAX_RENDER_POINTS, + ), + [lrHistory], + ); + + const steps = useMemo(() => { + const set = new Set(); + for (const p of lossData) set.add(p.step); + for (const p of lrData) set.add(p.step); + return Array.from(set).sort((a, b) => a - b); + }, [lossData, lrData]); + + const stepDomain = useMemo(() => fullStepDomain(steps), [steps]); + const xAxisTicks = useMemo( + () => buildStepTicks(stepDomain[0], stepDomain[1]), + [stepDomain], + ); + + const lossDomain = useMemo( + () => buildYDomain(lossData.flatMap((p) => [p.displayLoss, p.displaySmoothed])), + [lossData], + ); + const lrDomain = useMemo( + () => buildYDomain(lrData.map((p) => p.displayLr)), + [lrData], + ); + + const avgRaw = + lossItems.length > 0 + ? +(lossItems.reduce((s, p) => s + p.loss, 0) / lossItems.length).toFixed(4) + : 0; + + if (lossItems.length === 0 && lrData.length === 0) return null; + + return ( +
+
+ +

+ Per-step loss is noisy by design: every step samples a random noise level. Watch + the smoothed line for the trend, not the raw jitter. +

+
+ +
+ ); +} 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 43ed22a234..ee97811214 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -6,10 +6,6 @@ import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; -// Deep import (not the studio feature index) so the Images bundle does not statically pull -// in the heavy StudioPage; charts-section itself lazy-loads its recharts content. -// eslint-disable-next-line no-restricted-imports -import { ChartsSection } from "@/features/studio/sections/charts-section"; import type { TrainingSeriesPoint } from "@/features/training"; // eslint-disable-next-line no-restricted-imports -- matches images-page.tsx's token access import { getHfToken, hfApiToken } from "@/features/hub/stores/hf-token-store"; @@ -27,6 +23,8 @@ import { uploadDiffusionDataset, } from "../api"; import { DatasetLabelingGrid, LabelingGridToggle } from "./dataset-labeling-grid"; +import { DatasetShowcase } from "./dataset-showcase"; +import { DiffusionCharts } from "./diffusion-charts"; import { ExampleDatasetCards } from "./example-dataset-cards"; // The families the Train tab can train, in the popularity order the user asked for. This is @@ -549,6 +547,14 @@ export function DiffusionTrainPanel({ ) : ( selectedDataset && ( <> + {selectedDataset.image_count > 0 && !gridOpen && ( + setGridOpen(true)} + /> + )} - + {completed && (