From 1601b7828fbe9e3cef3003d9311ac3ace297948d Mon Sep 17 00:00:00 2001
From: Daniel Han-Chen <danielhanchen@gmail.com>
Date: Mon, 25 May 2026 00:20:32 +0000
Subject: [PATCH] Fix safetensors chat backend unload for PR #5754

_release_chat_backend_for_diffusion was importing
get_inference_backend from core.inference.inference (the in-subprocess
class) and calling unload_model() without the required model_name
argument. The TypeError was swallowed and the active chat model
stayed resident, defeating the chat-to-diffusion lifecycle handoff.

Switch to the orchestrator's accessor at core.inference and pass
active_model_name through, mirroring the GGUF chat-load path. Add a
regression test that stubs both backends and verifies unload_model
is called with the active model name.
---
 studio/backend/core/inference/diffusion.py    | 18 ++++++---
 .../backend/tests/test_diffusion_backend.py   | 37 +++++++++++++++++++
 2 files changed, 49 insertions(+), 6 deletions(-)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index 50bebdd98c..aa1610a450 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -623,17 +623,23 @@ def _release_chat_backend_for_diffusion() -> None:
     except Exception as exc:
         logger.debug("llama-server unload skipped: %s", exc)
 
-    # 2. Safetensors / HF chat backend (the Inference orchestrator that
+    # 2. Safetensors / HF chat backend (the InferenceOrchestrator that
     #    serves FastVisionModel / FastLanguageModel weights). When this
     #    backend has a model resident on the same GPU, a diffusion load
-    #    will OOM the same way.
+    #    will OOM the same way. The orchestrator's unload_model takes a
+    #    model_name; passing it without args raised TypeError and was
+    #    swallowed, leaving the chat model resident.
     try:
-        from core.inference.inference import get_inference_backend  # type: ignore
+        from core.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()
+        active_model_name = getattr(backend, "active_model_name", None)
+        if active_model_name:
+            logger.info(
+                "Unloading safetensors chat backend '%s' before diffusion load",
+                active_model_name,
+            )
+            backend.unload_model(active_model_name)
     except Exception as exc:
         logger.debug("safetensors unload skipped: %s", exc)
 
diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py
index b3cf13f10d..56aa7255b8 100644
--- a/studio/backend/tests/test_diffusion_backend.py
+++ b/studio/backend/tests/test_diffusion_backend.py
@@ -477,6 +477,43 @@ def test_smart_base_repo_picks_base_4b(monkeypatch):
     assert status["base_repo"] == "black-forest-labs/FLUX.2-klein-base-4B"
 
 
+def test_release_chat_backend_calls_unload_with_model_name(monkeypatch):
+    """The safetensors backend unload helper must call unload_model
+    with the active model name (the orchestrator's signature requires
+    it). The previous behaviour swallowed TypeError and left the chat
+    model resident, defeating the lifecycle handoff."""
+    import sys
+    import types
+
+    fake_pkg = types.ModuleType("core.inference")
+    calls: list = []
+
+    class _Stub:
+        active_model_name = "owner/some-model"
+
+        def unload_model(self, name):
+            calls.append(name)
+            self.active_model_name = None
+            return True
+
+    stub = _Stub()
+    fake_pkg.get_inference_backend = lambda: stub
+    monkeypatch.setitem(sys.modules, "core.inference", fake_pkg)
+
+    # Skip the llama-server branch by also stubbing routes.inference.
+    fake_routes = types.ModuleType("routes.inference")
+    fake_routes.get_llama_cpp_backend = lambda: types.SimpleNamespace(
+        is_loaded = False
+    )
+    monkeypatch.setitem(sys.modules, "routes.inference", fake_routes)
+
+    from core.inference.diffusion import _release_chat_backend_for_diffusion
+
+    _release_chat_backend_for_diffusion()
+    assert calls == ["owner/some-model"], calls
+    assert stub.active_model_name is None
+
+
 def test_load_model_uses_safetensors_flag(monkeypatch):
     """The pipeline.from_pretrained call must pass use_safetensors=True
     so pickle-backed .bin weights are refused at load time."""