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

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2026-06-25 04:11:46 +00:00
commit 39eb022b52
2 changed files with 27 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -130,7 +130,12 @@ class DiffusionBackend:
return hf_hub_download(repo_id, gguf_filename, token = hf_token)
def _prefetch_files(
self, repo_id: str, gguf_filename: Optional[str], base: str, base_files: list[str], hf_token: Optional[str]
self,
repo_id: str,
gguf_filename: Optional[str],
base: str,
base_files: list[str],
hf_token: Optional[str],
) -> None:
"""Pre-download the GGUF + the given ``base_files`` into the HF cache,
WITHOUT the lock and honoring ``_cancel_event``, so load_pipeline's
@ -140,12 +145,16 @@ class DiffusionBackend:
# GGUF transformer (hub repos only; a local path is already on disk).
if gguf_filename and not Path(repo_id).expanduser().exists():
hf_hub_download_with_xet_fallback(repo_id, gguf_filename, hf_token, cancel_event = self._cancel_event)
hf_hub_download_with_xet_fallback(
repo_id, gguf_filename, hf_token, cancel_event = self._cancel_event
)
# Base repo (VAE / text-encoder / scheduler); list comes from the estimate.
for rfilename in base_files:
if self._cancel_event.is_set():
raise RuntimeError("Cancelled")
hf_hub_download_with_xet_fallback(base, rfilename, hf_token, cancel_event = self._cancel_event)
hf_hub_download_with_xet_fallback(
base, rfilename, hf_token, cancel_event = self._cancel_event
)
# ── Background load + progress ─────────────────────────────────────────
@ -219,7 +228,11 @@ class DiffusionBackend:
# Download outside the lock so unload()/an eviction can preempt the
# multi-GB pull; load_pipeline below then assembles from the cache.
self._prefetch_files(
kwargs["repo_id"], kwargs.get("gguf_filename"), base, base_files, kwargs.get("hf_token")
kwargs["repo_id"],
kwargs.get("gguf_filename"),
base,
base_files,
kwargs.get("hf_token"),
)
self.load_pipeline(**kwargs)
with self._lock:

14
studio/backend/tests/test_diffusion_backend.py
View file

@ -465,8 +465,11 @@ def test_prefetch_aborts_when_cancelled(tmp_path):
(tmp_path / "model.gguf").write_bytes(b"x")
with pytest.raises(RuntimeError, match = "Cancelled"):
backend._prefetch_files(
str(tmp_path), "model.gguf", "Tongyi-MAI/Z-Image-Turbo",
["vae/diffusion_pytorch_model.safetensors"], None,
str(tmp_path),
"model.gguf",
"Tongyi-MAI/Z-Image-Turbo",
["vae/diffusion_pytorch_model.safetensors"],
None,
)
@ -479,8 +482,11 @@ def test_prefetch_downloads_gguf_and_base(monkeypatch, tmp_path):
)
# Hub repo: the GGUF transformer and each base file are fetched.
backend._prefetch_files(
"unsloth/Z-Image-Turbo-GGUF", "model.gguf", "base/repo",
["vae/x.safetensors", "text_encoder/y.safetensors"], "hf_tok",
"unsloth/Z-Image-Turbo-GGUF",
"model.gguf",
"base/repo",
["vae/x.safetensors", "text_encoder/y.safetensors"],
"hf_tok",
)
assert ("unsloth/Z-Image-Turbo-GGUF", "model.gguf") in calls
assert ("base/repo", "vae/x.safetensors") in calls