Coordinate diffusion with training VRAM and guard its cache deletes

studio/backend/routes/inference.py

@@ -10178,6 +10178,29 @@ async def _openai_passthrough_non_streaming(
 # ──────────────────────────────────────────────────────────────────────────
 
 
+def _guard_diffusion_load_against_training() -> None:
+    """Refuse loading an image model while a training run is active. Unlike chat,
+    a diffusion pipeline's VRAM can't be cheaply estimated before the load, so the
+    load is refused outright rather than fit-checked. No-op when training is
+    inactive or its state can't be read. Raises HTTP 409."""
+    from core.training import get_training_backend
+
+    try:
+        if not get_training_backend().is_training_active():
+            return
+    except Exception as e:
+        logger.warning("Could not check training state for image-load guard: %s", e)
+        return
+    raise HTTPException(
+        status_code = 409,
+        detail = (
+            "Can't load an image model while training is running: the diffusion "
+            "pipeline would compete with the training run for GPU memory. Training "
+            "was left untouched. Try again after training finishes."
+        ),
+    )
+
+
 @studio_router.post("/images/load", response_model = DiffusionStatusResponse)
 async def load_diffusion_model(
     request: DiffusionLoadRequest, current_subject: str = Depends(get_current_subject)
@@ -10196,6 +10219,10 @@ async def load_diffusion_model(
             gguf_filename = request.gguf_filename,
             family_override = request.family_override,
         )
+        # Refuse while training is running: a multi-GB diffusion pipeline would
+        # compete with the training subprocess for VRAM. The chat path does the
+        # same via _guard_chat_load_against_training; this is its image sibling.
+        _guard_diffusion_load_against_training()
         # Take the GPU from the chat backend, then kick the (slow) load onto a
         # background thread and return at once — the client polls images/load-progress.
         await asyncio.to_thread(acquire_for, DIFFUSION)

studio/backend/routes/models.py

@@ -3331,6 +3331,24 @@ async def delete_cached_model(
     except Exception:
         pass
 
+    # Also refuse if the diffusion (Images) backend has this repo loaded; its
+    # delete guard is otherwise chat-only, so its GGUF could be removed from
+    # under a live pipeline. Repo-level match, like the chat guards above.
+    try:
+        from core.inference.diffusion import get_diffusion_backend
+        diffusion_status = get_diffusion_backend().status()
+        if diffusion_status.get("loaded") and diffusion_status.get("repo_id"):
+            loaded_id = str(diffusion_status["repo_id"]).lower()
+            if loaded_id == repo_id.lower() or loaded_id.startswith(repo_id.lower()):
+                raise HTTPException(
+                    status_code = 400,
+                    detail = "Unload the model before deleting",
+                )
+    except HTTPException:
+        raise
+    except Exception:
+        pass
+
     try:
         cache_scans = _all_hf_cache_scans()
 

studio/backend/routes/training.py

@@ -362,6 +362,25 @@ async def start_training(
             except Exception as e:
                 logger.warning("Could not shut down export subprocess: %s", e)
 
+            try:
+                # A resident or in-flight diffusion (Images) pipeline also holds
+                # GPU memory the training run needs, and it can't be cheaply sized,
+                # so tear it down unconditionally like the export subprocess above
+                # (the chat block below fit-checks; diffusion can't). unload() is a
+                # no-op when nothing is loaded and also preempts an in-flight load;
+                # release the arbiter so it doesn't think the gone pipeline owns
+                # the GPU. Must precede the chat block, which early-returns.
+                from core.inference import gpu_arbiter
+                from core.inference.diffusion import get_diffusion_backend
+
+                diffusion = get_diffusion_backend()
+                if diffusion.is_loaded:
+                    logger.info("Unloading diffusion (Images) model to free GPU memory for training")
+                diffusion.unload()
+                gpu_arbiter.release(gpu_arbiter.DIFFUSION)
+            except Exception as e:
+                logger.warning("Could not unload diffusion model for training: %s", e)
+
             try:
                 from routes.training_vram import (
                     can_keep_chat_during_training,

studio/backend/tests/test_cached_gguf_routes.py

@@ -706,3 +706,39 @@ def test_gguf_download_progress_counts_quant_subdir(monkeypatch, tmp_path):
 
     assert result["downloaded_bytes"] == 20_000
     assert result["progress"] == 1.0
+
+
+def test_delete_cached_refuses_diffusion_loaded_repo(monkeypatch):
+    # The cached-delete guard refuses deleting a repo the diffusion (Images)
+    # backend has loaded, mirroring the chat guard, so its GGUF can't be removed
+    # from under a live pipeline.
+    from fastapi import HTTPException
+    import core.inference.diffusion as diffusion_mod
+    import routes.inference as routes_inference
+
+    # Chat and orchestrator report nothing loaded; only diffusion holds the repo.
+    # delete_cached_model resolves get_inference_backend from the models module
+    # namespace, so patch it there (not on core.inference) to isolate that guard.
+    monkeypatch.setattr(
+        routes_inference, "get_llama_cpp_backend",
+        lambda: SimpleNamespace(is_loaded = False, model_identifier = None),
+    )
+    monkeypatch.setattr(
+        models_route, "get_inference_backend",
+        lambda: SimpleNamespace(active_model_name = None),
+    )
+    monkeypatch.setattr(
+        diffusion_mod, "get_diffusion_backend",
+        lambda: SimpleNamespace(status = lambda: {"loaded": True, "repo_id": "org/Z-Image-GGUF"}),
+    )
+
+    try:
+        asyncio.run(
+            models_route.delete_cached_model(
+                repo_id = "org/Z-Image-GGUF", variant = None, current_subject = "u",
+            )
+        )
+        assert False, "expected HTTPException refusing the delete"
+    except HTTPException as e:
+        assert e.status_code == 400
+        assert "Unload the model before deleting" in e.detail

studio/backend/tests/test_diffusion_routes.py

@@ -282,6 +282,31 @@ def test_load_validation_failure_does_not_evict_chat(client, monkeypatch):
     assert gpu_arbiter.current_owner() == gpu_arbiter.CHAT
 
 
+def test_load_refused_during_training_does_not_evict_chat(client, monkeypatch):
+    # An image load while training is active is refused (409) before the GPU is
+    # taken, so the training run and the loaded chat model are both untouched.
+    import core.training as core_training
+
+    monkeypatch.setattr(gpu_arbiter, "_owner", gpu_arbiter.CHAT)
+    evicted = []
+    monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.CHAT, lambda: evicted.append(True))
+
+    class _Training:
+        def is_training_active(self):
+            return True
+
+    monkeypatch.setattr(core_training, "get_training_backend", lambda: _Training())
+
+    resp = client.post(
+        "/api/inference/images/load",
+        json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"},
+    )
+    assert resp.status_code == 409
+    assert "training" in resp.json()["detail"].lower()
+    assert evicted == []  # chat backend was never evicted
+    assert gpu_arbiter.current_owner() == gpu_arbiter.CHAT
+
+
 def test_load_progress_route(client):
     # Before load: idle.
     idle = client.get("/api/inference/images/load-progress")