From 15d1de6e05b7d91f298f4823f60160dcf232be6f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 20:55:57 +0000 Subject: [PATCH] Show only loss and learning-rate charts for diffusion training The Train tab reused the LLM charts section, which also rendered an empty Grad Norm card and an Eval Loss card showing an Evaluation not configured placeholder with a red smear. Neither applies to diffusion LoRA training. Add a diffusion-only two-card view that reuses the loss and learning-rate cards directly with fixed presentation defaults, and note under the loss chart that per-step loss is noisy by design so users read the smoothed line for the trend rather than the raw jitter. --- .../images/train/diffusion-charts.tsx | 143 ++++++++++++++++++ .../images/train/diffusion-train-panel.tsx | 25 ++- 2 files changed, 154 insertions(+), 14 deletions(-) create mode 100644 studio/frontend/src/features/images/train/diffusion-charts.tsx 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 && (