Fix/adjust diffusion symmetric chat handoff for PR #5754

_release_chat_backend_for_diffusion now unloads both the GGUF
chat backend (llama-server) and the safetensors / HF chat backend
(get_inference_backend) before a diffusion load. Mirror the
behaviour on the chat-load side: both the Unsloth/transformers
load path and the GGUF load path now unload the diffusion pipeline
before claiming GPU memory. Closes the OOM-on-swap path flagged
by reviewers in both directions.

studio/backend/core/inference/diffusion.py

@@ -535,30 +535,40 @@ def _release_chat_backend_for_diffusion() -> None:
     """Unload any running chat backend before a diffusion load.
 
     Diffusion pipelines on FLUX-class models can eat 12-24 GB of VRAM,
-    and llama-server typically holds onto its loaded GGUF until told to
-    drop it. Asking the chat backend to release its weights first means
-    a typical 24 GB consumer GPU can host one chat model OR one
-    diffusion model without manual unload steps.
+    and the chat backends (llama-server for GGUF, the safetensors
+    Inference orchestrator for HF / Unsloth) typically hold onto their
+    loaded weights until told to drop them. Asking both to release
+    their weights first means a typical 24 GB consumer GPU can host
+    one chat model OR one diffusion model without manual unload steps.
 
-    Best effort: if the chat backend module is not importable (CI,
-    isolated tests, custom builds) we silently continue. Failures
-    inside the unload itself are logged but not propagated; the
-    diffusion load can still try and surface its own OOM.
+    Best effort: if a chat backend module is not importable (CI,
+    isolated tests, custom builds) or fails on the unload, we log and
+    continue; the diffusion load can still try and surface its own OOM.
     """
+    # 1. GGUF chat backend (llama-server subprocess).
     try:
         from routes.inference import get_llama_cpp_backend  # type: ignore
-    except Exception:
-        return
-    try:
+
         backend = get_llama_cpp_backend()
-    except Exception:
-        return
-    try:
         if getattr(backend, "is_loaded", False):
             logger.info("Unloading llama-server before diffusion load")
             backend.unload_model()
     except Exception as exc:
-        logger.warning("Could not unload chat backend before diffusion: %s", exc)
+        logger.debug("llama-server unload skipped: %s", exc)
+
+    # 2. Safetensors / HF chat backend (the Inference orchestrator that
+    #    serves FastVisionModel / FastLanguageModel weights). When this
+    #    backend has a model resident on the same GPU, a diffusion load
+    #    will OOM the same way.
+    try:
+        from core.inference.inference import get_inference_backend  # type: ignore
+
+        backend = get_inference_backend()
+        if getattr(backend, "active_model_name", None):
+            logger.info("Unloading safetensors chat backend before diffusion load")
+            backend.unload_model()
+    except Exception as exc:
+        logger.debug("safetensors unload skipped: %s", exc)
 
 
 def _release(obj: Any) -> None:

studio/backend/routes/inference.py

@@ -747,6 +747,19 @@ async def load_model(
                 )
                 unsloth_backend.unload_model(unsloth_backend.active_model_name)
 
+            # Symmetric with /images/load: drop any active diffusion
+            # pipeline so the GGUF chat load does not race the FLUX VAE
+            # for VRAM. Best effort; silently continue on failure.
+            try:
+                from core.inference.diffusion import get_diffusion_backend
+
+                diff_backend = get_diffusion_backend()
+                if diff_backend.is_loaded:
+                    logger.info("Unloading diffusion pipeline before GGUF load")
+                    diff_backend.unload_model()
+            except Exception as e:
+                logger.debug("diffusion unload skipped (GGUF path): %s", e)
+
             # Inherit llama_extra_args from the previous load when the
             # request omits the field (the chat-settings Apply path
             # does not round-trip them; explicit [] still clears).
@@ -923,6 +936,20 @@ async def load_model(
             logger.info("Unloading GGUF model before loading Unsloth model")
             llama_backend.unload_model()
 
+        # Unload any active diffusion pipeline so the new chat model is
+        # not racing the FLUX VAE for VRAM on a 16-24 GB card.
+        try:
+            from core.inference.diffusion import get_diffusion_backend
+
+            diff_backend = get_diffusion_backend()
+            if diff_backend.is_loaded:
+                logger.info(
+                    "Unloading diffusion pipeline before loading Unsloth chat model"
+                )
+                diff_backend.unload_model()
+        except Exception as e:
+            logger.debug("diffusion unload skipped: %s", e)
+
         # Shut down any export subprocess to free VRAM
         try:
             from core.export import get_export_backend