diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 226bcaa5ff..2239174690 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -369,6 +369,7 @@ def dit_accelerator_missing_reason(resolved_family: str) -> Optional[str]: return None try: import torch + xpu = getattr(torch, "xpu", None) mps = getattr(torch, "mps", None) if ( diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index e23eeda7ac..3d76543ee1 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -4614,9 +4614,11 @@ async def _load_model_impl( # requirement exactly as it does for the load. if config.is_gguf and config.gguf_hf_repo: from core.inference.llama_cpp import gguf_load_in_flight + gguf_load_stack.enter_context(gguf_load_in_flight(config.gguf_hf_repo)) from core.inference.llama_cpp import _hub_download_blocks_gguf_load + if await asyncio.to_thread( _hub_download_blocks_gguf_load, config.gguf_hf_repo, diff --git a/studio/backend/tests/test_diffusion_base_precision.py b/studio/backend/tests/test_diffusion_base_precision.py index 5316c79217..52a1318876 100644 --- a/studio/backend/tests/test_diffusion_base_precision.py +++ b/studio/backend/tests/test_diffusion_base_precision.py @@ -621,6 +621,7 @@ def test_assert_trusted_base_model_rejects_local_non_pipeline(tmp_path): with pytest.raises(ValueError, match = "untrusted"): common._assert_trusted_base_model("evil/base") + def test_dit_accelerator_missing_reason_and_info_hide_train_without_a_gpu(monkeypatch): # Clicking Start on a GPU-less host evicted the resident Images pipeline, downloaded the text # encoders, and only then died in the child: diffusers' bitsandbytes quantizer refuses 4-bit diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py index d453c8f220..b2eed1e19a 100644 --- a/studio/backend/tests/test_gpu_selection.py +++ b/studio/backend/tests/test_gpu_selection.py @@ -1151,9 +1151,7 @@ class TestRouteErrors(unittest.TestCase): patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), # The chat handoff passes a `register` hook (the in-flight marker), so accept it. - patch.object( - arb, "acquire_for", lambda owner, register = None: acquired.append(owner) - ), + patch.object(arb, "acquire_for", lambda owner, register = None: acquired.append(owner)), ): with self.assertRaises(HTTPException) as exc_info: asyncio.run(