diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 693fe964cd..34e0ec4746 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -57,7 +57,7 @@ _LR_SCHEDULERS: frozenset[str] = frozenset( # DiT families whose fp32 RoPE/embedder overflow fp16, so they train in bf16 only. Must stay # in sync with the DiT trainer's own specs (kept separate to avoid an import cycle). -_FORCE_BF16_FAMILIES: frozenset[str] = frozenset({"qwen-image", "z-image"}) +_FORCE_BF16_FAMILIES: frozenset[str] = frozenset({"qwen-image", "z-image", "krea-2"}) _IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"} _CAPTION_EXTS = (".txt", ".caption") diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index b35ac90692..0ef854cb70 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1609,14 +1609,22 @@ async def upload_diffusion_dataset( # an overwrite; caption/text files are exempt (sample.txt for sample.png is intended). if ext in _DIFFUSION_DATASET_IMAGE_EXTS: stem = Path(filename).stem + # Compare stems (and the same-name guard) case-insensitively: on Windows/macOS + # (case-insensitive filesystems) two images whose stems differ only by case + # (sample.png vs Sample.jpg) resolve to the SAME .txt caption sidecar, so a + # case-sensitive check would let both through and silently share -- and corrupt -- + # one caption. Casefolding the name guard too keeps a same-name case variant + # (sample.png vs Sample.png, one file / an overwrite on those filesystems) exempt. + stem_cf = stem.casefold() + fname_cf = filename.casefold() clash = next( ( p.name for p in folder.iterdir() if p.is_file() - and p.name != filename + and p.name.casefold() != fname_cf and p.suffix.lower() in _DIFFUSION_DATASET_IMAGE_EXTS - and p.stem == stem + and p.stem.casefold() == stem_cf ), None, ) @@ -1625,9 +1633,9 @@ async def upload_diffusion_dataset( ( n for n in names - if n != filename + if n.casefold() != fname_cf and Path(n).suffix.lower() in _DIFFUSION_DATASET_IMAGE_EXTS - and Path(n).stem == stem + and Path(n).stem.casefold() == stem_cf ), None, ) diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py index dd065b155e..92c73e4278 100644 --- a/studio/backend/tests/test_diffusion_lora_trainer.py +++ b/studio/backend/tests/test_diffusion_lora_trainer.py @@ -186,6 +186,30 @@ def test_config_normalized_lists_mxfp8_in_invalid_mode_error(): ).normalized() +def test_config_normalized_krea2_requires_bf16_compute(): + # krea-2 (like qwen-image / z-image) has fp32 RoPE/embedder internals that overflow fp16, + # so its DiT trains in bf16 only; fp16 must be refused up front by the route preflight, + # before it reserves training and evicts resident GPU models and the child trainer raises. + with pytest.raises(ValueError, match = "bf16"): + DiffusionLoraConfig( + base_model = "b", + data_dir = "d", + output_dir = "o", + model_family = "krea-2", + mixed_precision = "fp16", + ).normalized() + + +def test_force_bf16_families_matches_trainer_specs(): + # The route-level bf16-only preflight set must list exactly the DiT families whose trainer + # spec sets force_bf16. If a force_bf16 family is missing from the set (as krea-2 was), an + # fp16 start passes the route preflight, reserves training + evicts resident models, and + # only the child trainer raises -- the evict-then-fail the preflight exists to prevent. + from core.training.diffusion_dit_trainer import _SPECS + from core.training.diffusion_train_common import _FORCE_BF16_FAMILIES + assert _FORCE_BF16_FAMILIES == {fam for fam, spec in _SPECS.items() if spec.force_bf16} + + def _cfg(**kw): return DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", **kw) diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index 4fe30a8d58..c6b817470c 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -859,6 +859,49 @@ def test_diffusion_dataset_upload_normalizes_windows_and_rejects_dotdot(client, assert r.status_code == 400 and "Unsupported file" in r.json()["detail"] +def test_diffusion_dataset_upload_rejects_case_insensitive_stem_clash(client, dataset_roots): + # Two images whose stems differ only by case (sample.png vs Sample.jpg) map to the SAME + # .txt caption sidecar on case-insensitive filesystems (Windows / default macOS), so + # keeping both would silently share -- and corrupt -- one caption during training. The clash + # check compares stems with casefold, so it must reject the pair (the comparison is pure + # string logic, so this fires regardless of the test host filesystem). + r = client.post( + "/api/train/diffusion/dataset", + data = {"name": "caseset"}, + files = [ + ("files", ("sample.png", b"p", "image/png")), + ("files", ("Sample.jpg", b"j", "image/jpeg")), + ], + ) + assert r.status_code == 400 and "Duplicate image name" in r.json()["detail"] + + # The same clash across batches (the new image collides with one already on disk). + r = client.post( + "/api/train/diffusion/dataset", + data = {"name": "caseset2"}, + files = [("files", ("photo.png", b"p", "image/png"))], + ) + assert r.status_code == 200, r.text + r = client.post( + "/api/train/diffusion/dataset", + data = {"name": "caseset2"}, + files = [("files", ("PHOTO.webp", b"w", "image/webp"))], + ) + assert r.status_code == 400 and "Duplicate image name" in r.json()["detail"] + + # A same-name case variant with the SAME extension is an overwrite (one file on a + # case-insensitive FS), not a caption clash, so it is still allowed. + r = client.post( + "/api/train/diffusion/dataset", + data = {"name": "caseset3"}, + files = [ + ("files", ("pic.png", b"a", "image/png")), + ("files", ("Pic.png", b"b", "image/png")), + ], + ) + assert r.status_code == 200, r.text + + def test_diffusion_dataset_upload_over_cap_keeps_existing_example( client, dataset_roots, monkeypatch ):