Fix diffusion dataset 500s, the dropout-1.0 no-op run and the reset base pick

Four correctness fixes on the training side:

- The labeling grid read caption sidecars under except OSError, but a
  non-UTF-8 sidecar raises UnicodeDecodeError (a ValueError), so one bad
  file 500d /diffusion/dataset/{name}/images and the grid could not be
  opened to repair it. Read it as no caption, matching the info summary.
- An image past Pillow's own hard limit raises DecompressionBombError,
  which derives straight from Exception and so escaped the upload guard's
  (OSError, UnidentifiedImageError, ValueError) and returned 500 instead
  of the intended 400.
- lora_dropout accepted 1.0, which makes PEFT build nn.Dropout(p=1.0):
  lora_A and lora_B receive no gradient and the run saves an untrained
  adapter while reporting normal progress. Bound it below 1.0, matching
  the LLM request schema.
- The train panel re-seeded the base repo on every dataset refresh
  because the family object identity changes on each info fetch, so an
  upload or caption save silently replaced the user's chosen base and the
  run started on a different model. Track the pick and only re-seed on a
  real family change.

studio/backend/models/training.py

@@ -752,7 +752,10 @@ class DiffusionTrainingStartRequest(BaseModel):
     gradient_accumulation_steps: int = Field(1, ge = 1, le = 256)
     lora_rank: int = Field(16, ge = 1, le = 320)
     lora_alpha: Optional[int] = Field(None, ge = 1, le = 640, description = "Defaults to lora_rank")
-    lora_dropout: float = Field(0.0, ge = 0.0, le = 1.0)
+    # Strictly below 1: PEFT turns lora_dropout into nn.Dropout(p=...), so 1.0 zeroes every input
+    # to the LoRA branch -- lora_A/lora_B receive no gradient and the run saves an untrained
+    # adapter while reporting normal progress. Matches TrainingStartRequest's [0, 1) validator.
+    lora_dropout: float = Field(0.0, ge = 0.0, lt = 1.0)
     # Mirror the remaining training-affecting knobs of DiffusionLoraConfig so a client that sets
     # them isn't silently trained with defaults. Default targets to the SDXL attention
     # projections (the trainer's DEFAULT_LORA_TARGETS) so it is never None.

studio/backend/routes/training.py

@@ -1906,6 +1906,17 @@ def _validate_uploaded_training_image(path: Path, original_name: str) -> None:
     try:
         with Image.open(path) as image:
             width, height = image.size
+    except Image.DecompressionBombError:
+        # Past Pillow's own hard limit (> 2 x MAX_IMAGE_PIXELS ~ 179 MP) Image.open() raises before
+        # .size can be read. That error derives straight from Exception (not OSError/ValueError), so
+        # letting it escape 500s the upload; it is exactly the oversized image this guard rejects.
+        raise HTTPException(
+            status_code = 400,
+            detail = (
+                f"Image '{original_name}' is too large; maximum is "
+                f"{_MAX_TRAINING_IMAGE_SIDE}px per side."
+            ),
+        )
     except (OSError, UnidentifiedImageError, ValueError):
         return  # not a decodable image -> not a bomb; leave the existing contract
     if width > _MAX_TRAINING_IMAGE_SIDE or height > _MAX_TRAINING_IMAGE_SIDE:
@@ -1985,7 +1996,11 @@ def _image_record(
             try:
                 caption = sidecar.read_text(encoding = "utf-8").strip()
                 source = "sidecar"
-            except OSError:
+            except (OSError, UnicodeError):
+                # Unreadable / invalid UTF-8 sidecar (uploads store text sidecars as raw bytes):
+                # UnicodeDecodeError is a ValueError, not an OSError, so an OSError-only guard let
+                # it 500 the whole labeling grid. Read it as no caption, like the info summary's
+                # _resolve_dataset_caption does.
                 caption = None
             break
     if caption is None:

studio/backend/tests/test_diffusion_dataset_api.py

@@ -88,6 +88,34 @@ def test_list_images_caption_precedence(client, ds_root):
     assert recs["a.png"]["width"] == 8 and recs["a.png"]["height"] == 8
 
 
+def test_list_images_tolerates_invalid_utf8_sidecar(client, ds_root):
+    # The upload route stores .txt/.caption sidecars as raw bytes (and users hand-drop them), so a
+    # sidecar can hold non-UTF-8 text. read_text then raises UnicodeDecodeError, which is a
+    # ValueError -- NOT an OSError -- so an `except OSError` around it 500s the whole labeling grid
+    # and the user cannot open it to repair the caption. One bad sidecar must read as no caption
+    # while every other image still lists (the info summary already behaves this way).
+    folder = ds_root / "badutf8"
+    folder.mkdir()
+    _write_png(folder / "a.png")
+    _write_png(folder / "b.png")
+    (folder / "a.txt").write_bytes(b"\xff\xfe not valid utf-8")
+    (folder / "b.txt").write_text("cap b", encoding = "utf-8")
+
+    r = client.get("/api/train/diffusion/dataset/badutf8/images")
+    assert r.status_code == 200, r.text
+    recs = {rec["filename"]: rec for rec in r.json()["images"]}
+    assert set(recs) == {"a.png", "b.png"}
+    assert recs["a.png"]["caption"] in (None, "")
+    assert recs["b.png"]["caption"] == "cap b"
+
+    # The caption PUT returns the same record, so it must not 500 after writing either.
+    put = client.put(
+        "/api/train/diffusion/dataset/badutf8/caption/a.png", json = {"caption": "fixed"}
+    )
+    assert put.status_code == 200, put.text
+    assert put.json()["caption"] == "fixed"
+
+
 def test_list_images_missing_dataset_404(client, ds_root):
     assert client.get("/api/train/diffusion/dataset/nope/images").status_code == 404
 

studio/backend/tests/test_diffusion_training.py

@@ -588,6 +588,24 @@ def test_request_model_dit_loss_knob_bounds():
             DiffusionTrainingStartRequest(**base, **bad)
 
 
+def test_request_model_rejects_lora_dropout_of_one():
+    # lora_dropout = 1.0 makes PEFT build nn.Dropout(p=1.0), which zeroes every input to the LoRA
+    # branch: the adapter output is identically the frozen base and both lora_A/lora_B get exactly
+    # zero gradient, so the run reports normal progress and saves an untrained adapter. The generic
+    # training schema already requires < 1 (TrainingStartRequest._check_lora_dropout); the diffusion
+    # schema must match instead of accepting the degenerate boundary.
+    from pydantic import ValidationError
+
+    from models.training import DiffusionTrainingStartRequest
+
+    base = {"base_model": "b", "data_dir": "d", "output_dir": "o"}
+    assert DiffusionTrainingStartRequest(**base).lora_dropout == 0.0
+    assert DiffusionTrainingStartRequest(**base, lora_dropout = 0.99).lora_dropout == 0.99
+    for bad in (1.0, 1.5, -0.1):
+        with pytest.raises(ValidationError):
+            DiffusionTrainingStartRequest(**base, lora_dropout = bad)
+
+
 def test_request_model_num_epochs_bounds():
     # The request schema mirrors DiffusionLoraConfig's 0..1000 num_epochs range.
     from pydantic import ValidationError
@@ -1109,6 +1127,27 @@ def test_diffusion_dataset_upload_rejects_oversized_image(client, dataset_roots)
     assert r2.status_code == 200, r2.text
 
 
+def test_diffusion_dataset_upload_maps_pillow_bomb_to_400(client, dataset_roots, monkeypatch):
+    # Past Pillow's hard bomb threshold (> 2 x Image.MAX_IMAGE_PIXELS) Image.open() itself raises
+    # DecompressionBombError, which derives straight from Exception -- not OSError/ValueError -- so
+    # the dimension guard's except clause missed it and the upload 500'd instead of returning the
+    # intended oversized-image 400. Shrink the limit (as Pillow's own tests do) so a small file
+    # crosses it, instead of building a 179 MP image in the test.
+    pytest.importorskip("PIL")
+    from PIL import Image
+
+    monkeypatch.setattr(Image, "MAX_IMAGE_PIXELS", 8)  # 8x8 = 64 pixels > 2 x 8
+    r = client.post(
+        "/api/train/diffusion/dataset",
+        data = {"name": "bomb-hard"},
+        files = [("files", ("huge.png", _png_bytes(8, 8), "image/png"))],
+    )
+    assert r.status_code == 400, r.text
+    assert "too large" in r.json()["detail"]
+    # All-or-nothing: the rejected image is not left in the dataset.
+    assert not (dataset_roots[0] / "bomb-hard" / "huge.png").exists()
+
+
 def test_diffusion_info_tolerates_non_object_jsonl(client, dataset_roots):
     # A metadata.jsonl line that is valid JSON but not an object ([]/null/string/number) or malformed
     # must be skipped per-line, not 500 the info endpoint; a valid row in the same file still counts.

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

@@ -297,6 +297,14 @@ export function DiffusionTrainPanel({
   // Track whether the user hand-picked a base precision; if not, a family change re-seeds it
   // from that family's recommended_precision.
   const precisionDirty = useRef(false);
+  // Same for the base repo: once the user picks one, only a real family change may re-seed it.
+  // `family` is a fresh object after every refreshInfo() (mergeFamilies rebuilds the list from the
+  // new info), so the seeding effect below re-runs on an unrelated dataset refresh (upload, import,
+  // caption save) too; without this the user's chosen base was silently replaced by the family
+  // default and the run started on a different model. The family the base was last seeded for is
+  // tracked by name, since the object identity is not stable.
+  const baseDirty = useRef(false);
+  const seededBaseFamily = useRef<string | null>(null);
 
   const [starting, setStarting] = useState(false);
   const [status, setStatus] = useState<DiffusionTrainingStatus | null>(null);
@@ -400,11 +408,17 @@ export function DiffusionTrainPanel({
   // user edited the numbers). Prefer the loaded base repo when it belongs to this family.
   useEffect(() => {
     if (!family) return;
+    // A NEW family invalidates any earlier base pick (another family's repo is not selectable
+    // here); a mere info refresh does not, so compare by name rather than object identity.
+    if (seededBaseFamily.current !== family.name) {
+      seededBaseFamily.current = family.name;
+      baseDirty.current = false;
+    }
     const preferLoaded =
       loadedBaseRepo && family.base_repos.includes(loadedBaseRepo)
         ? loadedBaseRepo
         : family.base_repos[0] ?? CUSTOM_BASE;
-    setBaseChoice(preferLoaded);
+    if (!baseDirty.current) setBaseChoice(preferLoaded);
     if (!settingsDirty.current) {
       setLearningRate(family.defaults.lr);
       setRank(family.defaults.rank);
@@ -1063,7 +1077,13 @@ export function DiffusionTrainPanel({
 
           <div className={fieldClass}>
             <Label className="text-xs">Base model</Label>
-            <Select value={effectiveBase} onValueChange={setBaseChoice}>
+            <Select
+              value={effectiveBase}
+              onValueChange={(v) => {
+                baseDirty.current = true; // an explicit pick survives later info refreshes
+                setBaseChoice(v);
+              }}
+            >
               <SelectTrigger className={selectClass} aria-label="Base model">
                 <SelectValue />
               </SelectTrigger>