Merge remote-tracking branch 'origin/diffusion-lora-training-api' into diffusion-lora-training-ui

studio/backend/routes/training.py

@@ -1158,6 +1158,15 @@ async def start_diffusion_training(
     except ValueError as e:
         raise HTTPException(status_code = 400, detail = str(e))
 
+    # Validate the config BEFORE freeing resident GPU workloads, so a start that is
+    # then refused (bad numbers, a non-SDXL base model) never tears down the user's
+    # loaded chat/Images model. service.start() re-runs this cheaply before spawn.
+    from core.training.diffusion_lora_trainer import _config_from_dict
+    try:
+        _config_from_dict(config).normalized()
+    except ValueError as e:
+        raise HTTPException(status_code = 400, detail = str(e))
+
     # Free resident GPU workloads (export / Images pipeline / chat) before the trainer
     # loads its own SDXL pipeline.
     _free_gpu_for_diffusion_training()
@@ -1253,6 +1262,7 @@ _DATASET_NAME_RE = None  # compiled lazily; module keeps its import block torch-
 def _clean_diffusion_dataset_name(name: str) -> str:
     """Validate a dataset folder name: a single path component, no traversal, printable."""
     import re
+
     global _DATASET_NAME_RE
     if _DATASET_NAME_RE is None:
         _DATASET_NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._ -]{0,127}$")

studio/backend/tests/test_diffusion_training.py

@@ -459,3 +459,19 @@ def test_diffusion_dataset_upload_rejects_unsupported_files(client, dataset_root
     )
     assert r.status_code == 400
     assert "Unsupported file" in r.json()["detail"]
+
+
+def test_route_start_refuses_non_sdxl_base_without_freeing_gpu(client, monkeypatch):
+    # A doomed start (non-SDXL base) must 400 BEFORE resident GPU workloads are freed,
+    # so a bad pick never unloads the user's working chat/Images model.
+    import routes.training as tr
+
+    freed = []
+    monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: freed.append(1))
+    r = client.post(
+        "/api/train/diffusion/start", json = {**_BODY, "base_model": "unsloth/FLUX.1-dev-GGUF"}
+    )
+    assert r.status_code == 400
+    assert "SDXL" in r.json()["detail"]
+    assert freed == []
+    assert client._fake.started_with is None