diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index e857ba3793..b9763bebed 100644 --- a/studio/backend/models/training.py +++ b/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. diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 3589abf604..3e0414ea61 100644 --- a/studio/backend/routes/training.py +++ b/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: diff --git a/studio/backend/tests/test_diffusion_dataset_api.py b/studio/backend/tests/test_diffusion_dataset_api.py index 641bf1ddad..ec7a580d95 100644 --- a/studio/backend/tests/test_diffusion_dataset_api.py +++ b/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 diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index cb45f0304b..b049ff8aa9 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/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. 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 d2081f7bd7..7c1f59765b 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/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(null); const [starting, setStarting] = useState(false); const [status, setStatus] = useState(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({
- { + baseDirty.current = true; // an explicit pick survives later info refreshes + setBaseChoice(v); + }} + >