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 a76e8e5209..813a6c1fd2 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -92,6 +92,21 @@ const FAMILY_PRESETS: FamilyPreset[] = [ const CUSTOM_BASE = "__custom__"; const UPLOAD_DATASET = "__upload__"; +// The dense DiT base precisions: they load a dense (bf16) base and quantise/cast it, so the +// backend rejects them for an already-quantised bnb-4bit repo. "nf4"/"auto" stay valid. +const DENSE_PRECISIONS = new Set(["bf16", "int8", "fp8"]); +// Mirror the backend's repo_is_prequantized heuristic: a repo whose name marks a +// bitsandbytes 4-bit build already ships a quantised transformer and cannot serve the dense +// base precisions. Kept in sync with diffusion_train_common.repo_is_prequantized. +function repoIsPrequantized(baseModel: string): boolean { + const name = baseModel.toLowerCase(); + return ( + name.includes("bnb-4bit") || + name.includes("-4bit") || + name.includes("int4") || + name.includes("nf4") + ); +} // 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"; @@ -372,6 +387,16 @@ export function DiffusionTrainPanel({ } }, [family, loadedBaseRepo, reportedFamily?.recommended_precision]); + // mixed_precision is an SDXL-only lever (its UI control is hidden for DiT families). A + // dense DiT base precision (bf16/int8/fp8) requires bf16 compute, and every DiT family + // trains in bf16, so reset precision to bf16 when the family changes to a DiT. Without + // this, an fp16/no value left over from SDXL rides along in the DiT start payload and the + // backend rejects it (dense modes need mixed_precision=bf16). Kept in its own effect so it + // does not re-trigger the base/settings reseed above. + useEffect(() => { + if (isDiT) setPrecision("bf16"); + }, [isDiT]); + // The base actually used everywhere (request, deploy, select value). baseChoice can // briefly hold another family's repo between a family switch and the reseed effect // (or if that effect is skipped); a raw