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

Merge remote-tracking branch 'origin/image-generation' into fix/imggen-review-bugs

Browse source
This commit is contained in:
oobabooga 2026-07-04 10:44:13 -03:00
commit 31ee7bae16
16 changed files with 297 additions and 56 deletions
Download patch file Download diff file Expand all files Collapse all files

1
studio/frontend/src/features/images/api.ts
View file

@ -291,6 +291,7 @@ export interface DiffusionMetricHistory {
steps: number[];
loss: number[];
lr: Array<number | null>;
grad_norm: Array<number | null>;
}
// A snapshot of the current diffusion training job (GET /api/train/diffusion/status).

9
studio/frontend/src/features/images/images-page.tsx
View file

@ -481,7 +481,7 @@ function AdvancedSelect({
return (
<div className="flex flex-col gap-1">
<div className="flex items-center justify-between gap-2">
<span className="flex items-center gap-1 text-xs font-medium text-muted-foreground">
<span className="flex shrink-0 items-center gap-1 whitespace-nowrap text-xs font-medium text-muted-foreground">
{label}
{hint && <InfoHint>{hint}</InfoHint>}
</span>
@ -1899,8 +1899,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
to GGUF (or nothing loaded) and otherwise show why it is unavailable. */}
{!status?.loaded || status.model_kind === "gguf" ? (
<AdvancedSelect
label="GGUF compute"
desc="Off runs the GGUF as-is. INT8/FP8/FP4 dequantise the transformer onto low-precision tensor cores for a faster step, at the cost of a larger download and more VRAM."
label="Dtype"
hint="Optional speed-up for GGUF models. Off runs the GGUF as-is. FP8/INT8/FP4 instead load the FULL base model and quantise its transformer onto low-precision tensor cores: faster per step, but a larger download and more VRAM, and it falls back to the GGUF if it can't fit. Needs CUDA."
value={transformerQuant}
onValueChange={(v) => setTransformerQuant(v as typeof transformerQuant)}
@ -1916,7 +1915,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
) : (
<div className="flex items-center justify-between gap-2">
<span className="flex items-center gap-1 text-xs font-medium text-muted-foreground">
GGUF compute
Dtype
</span>
<span className="text-xs text-muted-foreground/60">GGUF models only</span>
</div>
@ -2624,7 +2623,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
<p className="text-sm">
{status?.loaded
? "Enter a prompt and hit Generate."
: "Select a model quant to load, then generate."}
: "Select a diffusion model to load"}
</p>
</div>
)}

36
studio/frontend/src/features/images/train/diffusion-charts.tsx
View file

@ -9,6 +9,8 @@ import type { TrainingSeriesPoint } from "@/features/training";
// 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 { GradNormChartCard } from "@/features/studio/sections/charts/grad-norm-chart-card";
// 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";
@ -43,14 +45,17 @@ function fullStepDomain(steps: number[]): [number, number] {
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.
// A diffusion-only metrics view: Training Loss and Learning Rate side by side, plus Grad
// Norm (the pre-clip total gradient norm; spikes flag instability that raw loss noise
// hides), with a note under the loss card explaining why per-step loss looks noisy.
export function DiffusionCharts({
lossHistory,
lrHistory,
gradNormHistory = [],
}: {
lossHistory: TrainingSeriesPoint[];
lrHistory: TrainingSeriesPoint[];
gradNormHistory?: TrainingSeriesPoint[];
}): ReactElement | null {
const lossItems = useMemo(() => toLossItems(lossHistory), [lossHistory]);
const smoothed = useMemo(
@ -82,12 +87,24 @@ export function DiffusionCharts({
[lrHistory],
);
const gradNormData = useMemo(
() =>
compressSeries(
gradNormHistory
.filter((p) => Number.isFinite(p.value))
.map((p) => ({ step: p.step, gradNorm: p.value, displayGradNorm: p.value })),
MAX_RENDER_POINTS,
),
[gradNormHistory],
);
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);
for (const p of gradNormData) set.add(p.step);
return Array.from(set).sort((a, b) => a - b);
}, [lossData, lrData]);
}, [lossData, lrData, gradNormData]);
const stepDomain = useMemo(() => fullStepDomain(steps), [steps]);
const xAxisTicks = useMemo(
@ -103,6 +120,10 @@ export function DiffusionCharts({
() => buildYDomain(lrData.map((p) => p.displayLr)),
[lrData],
);
const gradNormDomain = useMemo(
() => buildYDomain(gradNormData.map((p) => p.displayGradNorm)),
[gradNormData],
);
const avgRaw =
lossItems.length > 0
@ -138,6 +159,15 @@ export function DiffusionCharts({
xAxisTicks={xAxisTicks}
scale="linear"
/>
{gradNormData.length > 0 && (
<GradNormChartCard
data={gradNormData}
domain={gradNormDomain}
visibleStepDomain={stepDomain}
xAxisTicks={xAxisTicks}
scale="linear"
/>
)}
</div>
);
}

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

@ -366,6 +366,13 @@ export function DiffusionTrainPanel({
.map((step, i) => ({ step, value: h.lr[i] }))
.filter((p): p is TrainingSeriesPoint => p.value != null);
}, [status?.metric_history]);
const gradNormHistory: TrainingSeriesPoint[] = useMemo(() => {
const h = status?.metric_history;
if (!h) return [];
return h.steps
.map((step, i) => ({ step, value: h.grad_norm?.[i] ?? null }))
.filter((p): p is TrainingSeriesPoint => p.value != null);
}, [status?.metric_history]);
const onUpload = useCallback(async () => {
const files = Array.from(fileInputRef.current?.files ?? []);
@ -791,7 +798,17 @@ export function DiffusionTrainPanel({
<>
<div className="bg-card corner-squircle flex flex-col gap-3 rounded-3xl p-5 ring-1 ring-foreground/10">
<div className="flex items-center justify-between">
<span className="text-sm font-semibold capitalize">{status.status}</span>
{/* A finished run should be unmistakable at a glance, so completed swaps
the plain status word for a celebratory line in the success color. */}
<span
className={
status.status === "completed"
? "text-sm font-semibold text-emerald-600 dark:text-emerald-400"
: "text-sm font-semibold capitalize"
}
>
{status.status === "completed" ? "Training complete \u{1F389}" : status.status}
</span>
<span className="text-xs text-muted-foreground">
{status.total_steps > 0 ? `${status.step}/${status.total_steps} steps` : ""}
</span>
@ -828,7 +845,11 @@ export function DiffusionTrainPanel({
)}
</div>
<DiffusionCharts lossHistory={lossHistory} lrHistory={lrHistory} />
<DiffusionCharts
lossHistory={lossHistory}
lrHistory={lrHistory}
gradNormHistory={gradNormHistory}
/>
{completed && (
<div className="bg-card corner-squircle flex flex-col gap-2 rounded-3xl p-5 ring-1 ring-foreground/10">