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.
This commit is contained in:
parent
7cd7c588f2
commit
15d1de6e05
2 changed files with 154 additions and 14 deletions
143
studio/frontend/src/features/images/train/diffusion-charts.tsx
Normal file
143
studio/frontend/src/features/images/train/diffusion-charts.tsx
Normal file
|
|
@ -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<number>();
|
||||
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 (
|
||||
<div className="grid grid-cols-1 gap-4 lg:grid-cols-2">
|
||||
<div className="flex flex-col gap-1">
|
||||
<TrainingLossChartCard
|
||||
data={lossData}
|
||||
domain={lossDomain}
|
||||
visibleStepDomain={stepDomain}
|
||||
xAxisTicks={xAxisTicks}
|
||||
avgRaw={avgRaw}
|
||||
avgDisplay={avgRaw}
|
||||
showRaw={true}
|
||||
showSmoothed={true}
|
||||
showAvgLine={true}
|
||||
scale="linear"
|
||||
/>
|
||||
<p className="px-1 text-[11px] leading-snug text-muted-foreground">
|
||||
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.
|
||||
</p>
|
||||
</div>
|
||||
<LearningRateChartCard
|
||||
data={lrData}
|
||||
domain={lrDomain}
|
||||
visibleStepDomain={stepDomain}
|
||||
xAxisTicks={xAxisTicks}
|
||||
scale="linear"
|
||||
/>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
|
@ -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 && (
|
||||
<DatasetShowcase
|
||||
dataset={dataset}
|
||||
imageCount={selectedDataset.image_count}
|
||||
refreshKey={gridRefresh}
|
||||
onBrowse={() => setGridOpen(true)}
|
||||
/>
|
||||
)}
|
||||
<LabelingGridToggle
|
||||
count={selectedDataset.image_count}
|
||||
open={gridOpen}
|
||||
|
|
@ -707,16 +713,7 @@ export function DiffusionTrainPanel({
|
|||
)}
|
||||
</div>
|
||||
|
||||
<ChartsSection
|
||||
currentStep={status.step}
|
||||
totalSteps={status.total_steps}
|
||||
isTraining={running}
|
||||
evalEnabled={false}
|
||||
lossHistory={lossHistory}
|
||||
lrHistory={lrHistory}
|
||||
gradNormHistory={[]}
|
||||
evalLossHistory={[]}
|
||||
/>
|
||||
<DiffusionCharts lossHistory={lossHistory} lrHistory={lrHistory} />
|
||||
|
||||
{completed && (
|
||||
<div className="bg-card corner-squircle flex flex-col gap-2 rounded-3xl p-5 ring-1 ring-foreground/10">
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue