ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix safetensors chat backend unload for PR #5754

Browse source 
_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.
This commit is contained in:
Daniel Han-Chen 2026-05-25 00:20:32 +00:00
commit 1601b7828f
2 changed files with 49 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

18
studio/backend/core/inference/diffusion.py
View file

@ -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)