From 373edae1c05ce0cc8b19a249b797db7f749fd398 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 2 Jul 2026 09:59:58 +0000
Subject: [PATCH] Validate diffusion training config before freeing the GPU

The start route freed resident GPU workloads (export, Images pipeline, chat)
before the service validated the config, so a start that was then refused,
now including a non-SDXL base model, tore down the user's loaded model for
nothing. Run the same cheap normalise pass first; the LLM path already
follows this rule via its before_spawn hook.
---
 studio/backend/routes/training.py               |  9 +++++++++
 studio/backend/tests/test_diffusion_training.py | 16 ++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py
index 63e3c74570..7d87f897ab 100644
--- a/studio/backend/routes/training.py
+++ b/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()
diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py
index 6fdfe7c009..9c4452df7f 100644
--- a/studio/backend/tests/test_diffusion_training.py
+++ b/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