Images Train LoRA dialog: token, validation, precision, base-repo prefill, gating, refresh

Nine review findings on the SDXL training dialog:
- Forward the saved Hub token so a gated/private SDXL base can be trained (the
  image load flow already sends it).
- Re-seed the base-model field from the current default each time the dialog
  opens; the keep-alive dialog otherwise kept its mount-time default after a
  model loaded.
- Prefill from base_repo (the diffusers pipeline) rather than repo_id, which for
  a GGUF/single-file SDXL load is the checkpoint path from_pretrained can't open.
- Add client-side validation of steps/rank/resolution/batch/learning-rate before
  the request.
- Expose a precision selector (bf16/fp16/fp32) so non-bf16 GPUs can train from
  the UI, not only the API.
- Gate the dialog on the active Images route (active && trainOpen) so switching
  tabs closes it and stops its polling.
- Rescan the LoRA picker when a run completes, so a freshly-trained adapter
  appears without a model reload.
- Cap the dialog height and scroll the body so the Start/Stop footer stays
  reachable on short viewports.
- Correct the copy to not over-promise picker auto-discovery.

Freeing the resident Images pipeline before training is handled backend-side in
the diffusion training start route.

studio/frontend/src/features/images/api.ts

@@ -269,10 +269,14 @@ export interface DiffusionTrainingStartRequest {
   gradient_accumulation_steps?: number;
   lora_rank?: number;
   lora_alpha?: number | null;
+  lora_target_modules?: string[];
+  max_grad_norm?: number;
   seed?: number;
   mixed_precision?: "bf16" | "fp16" | "no";
   gradient_checkpointing?: boolean;
   lr_scheduler?: string;
+  // Forwarded to StableDiffusionXLPipeline.from_pretrained for a gated/private base repo.
+  hf_token?: string | null;
 }
 
 // A snapshot of the current diffusion training job (GET /api/train/diffusion/status).

studio/frontend/src/features/images/diffusion-train-dialog.tsx

@@ -14,6 +14,7 @@ import {
 } from "@/components/ui/dialog";
 import { Input } from "@/components/ui/input";
 import { Label } from "@/components/ui/label";
+import { getHfToken, hfApiToken } from "@/features/hub/stores/hf-token-store";
 import { toast } from "@/lib/toast";
 
 import {
@@ -23,6 +24,8 @@ import {
   stopDiffusionTraining,
 } from "./api";
 
+const DEFAULT_SDXL_BASE = "stabilityai/stable-diffusion-xl-base-1.0";
+
 // A self-contained "Train a LoRA" dialog for the diffusion (SDXL) trainer. It posts to
 // /api/train/diffusion/start and polls /status while open, so it never blocks the page and
 // works whether or not a model is loaded for generation. Only SDXL is trainable today.
@@ -30,12 +33,15 @@ export function DiffusionTrainDialog({
   open,
   onOpenChange,
   defaultBaseModel,
+  onTrainingComplete,
 }: {
   open: boolean;
   onOpenChange: (open: boolean) => void;
   defaultBaseModel?: string;
+  // Called once when a run finishes so the page can rescan the LoRA picker.
+  onTrainingComplete?: () => void;
 }) {
-  const [baseModel, setBaseModel] = useState(defaultBaseModel || "stabilityai/stable-diffusion-xl-base-1.0");
+  const [baseModel, setBaseModel] = useState(defaultBaseModel || DEFAULT_SDXL_BASE);
   const [dataDir, setDataDir] = useState("");
   const [outputDir, setOutputDir] = useState("");
   const [instancePrompt, setInstancePrompt] = useState("");
@@ -44,9 +50,18 @@ export function DiffusionTrainDialog({
   const [rank, setRank] = useState(16);
   const [resolution, setResolution] = useState(1024);
   const [batchSize, setBatchSize] = useState(1);
+  const [precision, setPrecision] = useState<"bf16" | "fp16" | "no">("bf16");
   const [starting, setStarting] = useState(false);
   const [status, setStatus] = useState<DiffusionTrainingStatus | null>(null);
 
+  // The dialog stays mounted (ImagesPage is keep-alive), so the initial state seed does not
+  // reflect a base model loaded AFTER mount. Re-seed the base-model field from the current
+  // default each time the dialog opens, so "Train LoRA" after loading an SDXL checkpoint
+  // starts from that checkpoint's diffusers repo, not the hard-coded default.
+  useEffect(() => {
+    if (open) setBaseModel(defaultBaseModel || DEFAULT_SDXL_BASE);
+  }, [open, defaultBaseModel]);
+
   const poll = useCallback(async () => {
     try {
       setStatus(await getDiffusionTrainingStatus());
@@ -69,11 +84,33 @@ export function DiffusionTrainDialog({
       ? Math.min(100, Math.round((status.step / status.total_steps) * 100))
       : 0;
 
+  // Notify the parent exactly once when a run reaches "completed", so it can rescan the
+  // LoRA picker (a LoRA trained while a model is loaded is otherwise invisible until a
+  // model swap re-runs the discovery effect).
+  const [notifiedComplete, setNotifiedComplete] = useState(false);
+  useEffect(() => {
+    if (status?.status === "completed" && !notifiedComplete) {
+      setNotifiedComplete(true);
+      onTrainingComplete?.();
+    } else if (status?.status === "running" && notifiedComplete) {
+      setNotifiedComplete(false); // arm again for the next run
+    }
+  }, [status?.status, notifiedComplete, onTrainingComplete]);
+
   const onStart = useCallback(async () => {
     if (!baseModel.trim() || !dataDir.trim() || !outputDir.trim()) {
       toast.error("Base model, dataset folder, and output folder are required.");
       return;
     }
+    // Mirror the backend's numeric validation so obvious mistakes are caught before the
+    // request (the backend returns 400 for these; catching here gives a clearer message).
+    if (steps < 1) return toast.error("Steps must be at least 1.");
+    if (rank < 1) return toast.error("LoRA rank must be at least 1.");
+    if (resolution < 64 || resolution % 8 !== 0) {
+      return toast.error("Resolution must be a multiple of 8 and at least 64.");
+    }
+    if (batchSize < 1) return toast.error("Batch size must be at least 1.");
+    if (learningRate <= 0) return toast.error("Learning rate must be greater than 0.");
     setStarting(true);
     try {
       await startDiffusionTraining({
@@ -86,6 +123,10 @@ export function DiffusionTrainDialog({
         learning_rate: learningRate,
         train_batch_size: batchSize,
         lora_rank: rank,
+        mixed_precision: precision,
+        // Forward the saved Hub token so a gated/private SDXL base can be trained (the
+        // image load flow already sends it, so a model you can load, you can also train).
+        hf_token: hfApiToken(getHfToken()) || undefined,
       });
       toast.success("Training started");
       void poll();
@@ -94,7 +135,19 @@ export function DiffusionTrainDialog({
     } finally {
       setStarting(false);
     }
-  }, [baseModel, dataDir, outputDir, instancePrompt, resolution, steps, learningRate, batchSize, rank, poll]);
+  }, [
+    baseModel,
+    dataDir,
+    outputDir,
+    instancePrompt,
+    resolution,
+    steps,
+    learningRate,
+    batchSize,
+    rank,
+    precision,
+    poll,
+  ]);
 
   const onStop = useCallback(async () => {
     try {
@@ -108,17 +161,17 @@ export function DiffusionTrainDialog({
 
   return (
     <Dialog open={open} onOpenChange={onOpenChange}>
-      <DialogContent className="max-w-lg">
+      <DialogContent className="flex max-h-[85vh] max-w-lg flex-col">
         <DialogHeader>
           <DialogTitle>Train a LoRA (SDXL)</DialogTitle>
           <DialogDescription>
             Fine-tune an SDXL LoRA on a folder of images. Captions come from a metadata.jsonl,
-            per-image .txt sidecars, or the instance prompt below. The adapter is written to the
-            output folder and can be loaded from the LoRAs picker.
+            per-image .txt sidecars, or the instance prompt below. The adapter is saved to the
+            output folder shown after the run.
           </DialogDescription>
         </DialogHeader>
 
-        <div className="grid gap-3 py-2">
+        <div className="grid gap-3 overflow-y-auto py-2 pr-1">
           <div className="grid gap-1.5">
             <Label className="text-xs">Base model (SDXL repo or local path)</Label>
             <Input
@@ -200,16 +253,30 @@ export function DiffusionTrainDialog({
               />
             </div>
           </div>
-          <div className="grid gap-1.5">
-            <Label className="text-xs">Learning rate</Label>
-            <Input
-              type="number"
-              step={0.00001}
-              min={0}
-              value={learningRate}
-              onChange={(e) => setLearningRate(Number(e.target.value) || 0.0001)}
-              className="h-8 text-xs"
-            />
+          <div className="grid grid-cols-2 gap-3">
+            <div className="grid gap-1.5">
+              <Label className="text-xs">Learning rate</Label>
+              <Input
+                type="number"
+                step={0.00001}
+                min={0}
+                value={learningRate}
+                onChange={(e) => setLearningRate(Number(e.target.value) || 0.0001)}
+                className="h-8 text-xs"
+              />
+            </div>
+            <div className="grid gap-1.5">
+              <Label className="text-xs">Precision</Label>
+              <select
+                value={precision}
+                onChange={(e) => setPrecision(e.target.value as "bf16" | "fp16" | "no")}
+                className="h-8 rounded-md border border-input bg-background px-2 text-xs"
+              >
+                <option value="bf16">bf16 (default)</option>
+                <option value="fp16">fp16 (older GPUs)</option>
+                <option value="no">fp32 (no mixed)</option>
+              </select>
+            </div>
           </div>
 
           {status && status.status !== "idle" && (

studio/frontend/src/features/images/images-page.tsx

@@ -910,6 +910,9 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
   const [availableLoras, setAvailableLoras] = useState<DiffusionLoraInfo[]>([]);
   // "Train a LoRA" dialog (SDXL). Independent of the loaded generation model.
   const [trainOpen, setTrainOpen] = useState(false);
+  // Bumped when a training run completes, to force the LoRA discovery effect to rescan so
+  // a freshly-trained adapter appears in the picker without a model reload.
+  const [loraRefreshKey, setLoraRefreshKey] = useState(0);
   // ControlNet for the next generation: the chosen model id, a control image (data URL),
   // how to derive the control map, and the conditioning strength. Available models refresh
   // per loaded family; applied at generate time only when a model + control image are set.
@@ -1011,7 +1014,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
     return () => {
       cancelled = true;
     };
-  }, [loraCapable, status?.family]);
+  }, [loraCapable, status?.family, loraRefreshKey]);
 
   // Refresh the ControlNet picker's options when the loaded model (family) changes, and clear
   // a stale selection the new model can't use so an incompatible ControlNet is never sent.
@@ -1768,9 +1771,17 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
         </div>
       </div>
       <DiffusionTrainDialog
-        open={trainOpen}
+        open={active && trainOpen}
         onOpenChange={setTrainOpen}
-        defaultBaseModel={status?.family === "sdxl" ? status?.repo_id ?? undefined : undefined}
+        defaultBaseModel={
+          status?.family === "sdxl"
+            ? // Prefer base_repo (the full diffusers pipeline) over repo_id: for a GGUF or
+              // single-file SDXL load repo_id is the checkpoint path, which the trainer's
+              // from_pretrained cannot open. base_repo is the companion pipeline.
+              status?.base_repo ?? status?.repo_id ?? undefined
+            : undefined
+        }
+        onTrainingComplete={() => setLoraRefreshKey((k) => k + 1)}
       />
 
       {/* ── Controls rail + preview canvas. Padding mirrors the other tabs