diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 2d69c9bb56..ddf31db97a 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1582,7 +1582,14 @@ class DiffusionBackend: pipe = getattr(diffusers, pipe_cls_name).from_pipe( state.pipe, controlnet = cn_model, torch_dtype = None ) - self._cn_pipes[key] = pipe + with self._lock: + # Same race as the model cache above: an unload/superseding load may + # have cleared _cn_pipes while from_pipe ran; caching now would pin a + # pipeline built around the UNLOADED base and hand it to the next load. + if cancel.is_set() or self._state is not state: + del pipe + raise RuntimeError(DIFFUSION_CANCELLED_MSG) + self._cn_pipes[key] = pipe return pipe @staticmethod diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 75572b500a..baee6605e0 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -1096,6 +1096,13 @@ class SdCppDiffusionBackend: state.server.stop() if pending is not None and pending is not (state.server if state else None): pending.stop() + # Wait for a signalled one-shot generation to actually exit before reporting + # unloaded: callers (the GPU arbiter, training cleanup) treat this return as + # "the device is free", but a one-shot sd-cli child killed by the cancel above + # unwinds under _generate_lock. A bare acquire is the exit barrier (never taken + # while holding _lock; same pattern as DiffusionBackend.unload). + with self._generate_lock: + pass return self.status() def status(self) -> dict[str, Any]: diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index a89a8d01c2..56514d2272 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1234,6 +1234,21 @@ async def start_diffusion_training( # user's loaded chat/Images model, and never surfaces as a confusing mid-load 401. _preflight_gated_base(config.get("base_model", ""), config.get("hf_token")) + # Preflight the dataset too: a missing/empty/uncaptionable data_dir otherwise + # fails inside the spawned trainer AFTER the user's chat/Images model was + # evicted. Same discovery the trainer runs, so the two cannot disagree. + from core.training import diffusion_train_common as _dtc + + try: + await asyncio.to_thread( + _dtc.discover_image_caption_pairs, + config["data_dir"], + instance_prompt = config.get("instance_prompt") or None, + caption_column = config.get("caption_column") or "text", + ) + except (FileNotFoundError, 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 pipeline. _free_gpu_for_diffusion_training() @@ -1636,12 +1651,24 @@ async def set_diffusion_dataset_caption( sidecar = image_path.with_suffix(".txt") if caption: sidecar.write_text(caption, encoding = "utf-8") - else: - # Blank clears the sidecar; also drop a stale .caption so the image reads as - # uncaptioned afterwards. - sidecar.unlink(missing_ok = True) image_path.with_suffix(".caption").unlink(missing_ok = True) - return _image_record(folder, image_path, _load_metadata_captions(folder)) + return _image_record(folder, image_path, _load_metadata_captions(folder)) + # Blank must actually clear. Unlinking alone would resurface this image's + # metadata.jsonl / captions.jsonl caption (the fallback source), so when one + # exists write an EMPTY sidecar instead: both the record reader and the + # trainer's discovery treat an existing sidecar as authoritative even when + # empty, which makes it a tombstone. No metadata caption -> plain cleanup. + meta = _load_metadata_captions(folder) + try: + rel = image_path.relative_to(folder).as_posix() + except ValueError: + rel = image_path.name + if image_path.name in meta or rel in meta: + sidecar.write_text("", encoding = "utf-8") + else: + sidecar.unlink(missing_ok = True) + image_path.with_suffix(".caption").unlink(missing_ok = True) + return _image_record(folder, image_path, meta) return await asyncio.to_thread(write) diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py index 5708c3a3ae..3d4f7c98cf 100644 --- a/studio/backend/tests/test_diffusion_controlnet.py +++ b/studio/backend/tests/test_diffusion_controlnet.py @@ -239,6 +239,9 @@ def test_controlnet_pipe_loads_once_and_caches(monkeypatch): monkeypatch.setitem(sys.modules, "diffusers", _fake_diffusers()) b = DiffusionBackend() st = _state() + # The pipe cache only commits while ``st`` is the CURRENT load (an unload racing + # from_pipe must not repopulate the cache), so mirror the loaded invariant. + b._state = st resolved = dc.ResolvedControlNet("flux-union-pro", "repo/id", is_local = False) p1 = b._controlnet_pipe(st, resolved, threading.Event()) assert isinstance(p1, _FakeCNPipe) and isinstance(p1.controlnet, _FakeCNModel) @@ -259,3 +262,18 @@ def test_controlnet_pipe_rejects_family_without_classes(): st.family.controlnet_pipeline_class = None with pytest.raises(ValueError, match = "not supported"): b._controlnet_pipe(st, dc.ResolvedControlNet("x", "y", False), threading.Event()) + +def test_controlnet_pipe_not_cached_after_unload_race(monkeypatch): + # An unload that lands while from_pipe is assembling must not let the wrapper + # repopulate the cache around the torn-down base pipe. + import threading + + from core.inference.diffusion import DiffusionBackend + + monkeypatch.setitem(sys.modules, "diffusers", _fake_diffusers()) + b = DiffusionBackend() + st = _state() # never committed to b._state: the load is already gone + resolved = dc.ResolvedControlNet("flux-union-pro", "repo/id", is_local = False) + with pytest.raises(RuntimeError, match = "cancelled"): + b._controlnet_pipe(st, resolved, threading.Event()) + assert b._cn_pipes == {} diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index 1f85b08ccb..b5acaa7bb9 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -265,6 +265,13 @@ def client(monkeypatch): monkeypatch.setattr(tr, "get_training_backend", lambda: _FakeLLMBackend(active = False)) monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: None) + # The dataset preflight runs the trainer's discovery against _BODY's fake + # data_dir; stub it here so wiring tests pass, and let the dedicated preflight + # tests below re-point it at a real tmp dataset. + monkeypatch.setattr( + "core.training.diffusion_train_common.discover_image_caption_pairs", + lambda data_dir, **kw: [("img.png", "caption")], + ) app = FastAPI() app.include_router(training_router, prefix = "/api/train") app.dependency_overrides[get_current_subject] = lambda: "test-user"